parquet: add encrypted bloom filter fast path for known lengths

Author: fenfeng9

cpp/src/parquet/bloom_filter.cc

@@ -20,6 +20,7 @@
 #include <limits>
 #include <memory>
 
+#include "arrow/io/memory.h"
 #include "arrow/result.h"
 #include "arrow/util/logging_internal.h"
 #include "arrow/util/macros.h"
@@ -126,16 +127,10 @@ ::arrow::Status ValidateBloomFilterHeader(const format::BloomFilterHeader& heade
   return ::arrow::Status::OK();
 }
 
-}  // namespace
-
-BlockSplitBloomFilter BlockSplitBloomFilter::DeserializeEncrypted(
+BlockSplitBloomFilter DeserializeEncryptedFromStream(
     const ReaderProperties& properties, ArrowInputStream* input,
     std::optional<int64_t> bloom_filter_length, Decryptor* decryptor,
     int16_t row_group_ordinal, int16_t column_ordinal) {
-  if (decryptor == nullptr) {
-    throw ParquetException("Bloom filter decryptor must be provided");
-  }
-
   ThriftDeserializer deserializer(properties);
   format::BloomFilterHeader header;
 
@@ -158,6 +153,7 @@ BlockSplitBloomFilter BlockSplitBloomFilter::DeserializeEncrypted(
     throw ParquetException(
         "Bloom filter length less than encrypted bloom filter header length");
   }
+
   // Read the full header ciphertext and decrypt the Thrift header.
   auto header_cipher_buf =
       AllocateBuffer(properties.memory_pool(), header_cipher_total_len);
@@ -181,12 +177,17 @@ BlockSplitBloomFilter BlockSplitBloomFilter::DeserializeEncrypted(
   try {
     deserializer.DeserializeMessage(header_cipher_buf->data(), &header_cipher_len,
                                     &header, decryptor);
-    DCHECK_EQ(header_cipher_len, header_cipher_total_len);
   } catch (std::exception& e) {
     std::stringstream ss;
     ss << "Deserializing bloom filter header failed.\n" << e.what();
     throw ParquetException(ss.str());
   }
+  if (ARROW_PREDICT_FALSE(header_cipher_len != header_cipher_total_len)) {
+    std::stringstream ss;
+    ss << "Encrypted bloom filter header length mismatch: expected "
+       << header_cipher_total_len << " bytes, got " << header_cipher_len;
+    throw ParquetException(ss.str());
+  }
   PARQUET_THROW_NOT_OK(ValidateBloomFilterHeader(header));
 
   const int32_t bloom_filter_size = header.numBytes;
@@ -231,6 +232,34 @@ BlockSplitBloomFilter BlockSplitBloomFilter::DeserializeEncrypted(
   return bloom_filter;
 }
 
+}  // namespace
+
+BlockSplitBloomFilter BlockSplitBloomFilter::DeserializeEncrypted(
+    const ReaderProperties& properties, ArrowInputStream* input,
+    std::optional<int64_t> bloom_filter_length, Decryptor* decryptor,
+    int16_t row_group_ordinal, int16_t column_ordinal) {
+  if (decryptor == nullptr) {
+    throw ParquetException("Bloom filter decryptor must be provided");
+  }
+
+  // Read the full Bloom filter payload up front when the total length is known.
+  if (bloom_filter_length.has_value()) {
+    PARQUET_ASSIGN_OR_THROW(auto bloom_filter_buf, input->Read(*bloom_filter_length));
+    if (ARROW_PREDICT_FALSE(bloom_filter_buf->size() < *bloom_filter_length)) {
+      std::stringstream ss;
+      ss << "Bloom filter read failed: expected " << *bloom_filter_length
+         << " bytes, got " << bloom_filter_buf->size();
+      throw ParquetException(ss.str());
+    }
+    ::arrow::io::BufferReader reader(bloom_filter_buf);
+    return DeserializeEncryptedFromStream(properties, &reader, bloom_filter_length,
+                                          decryptor, row_group_ordinal, column_ordinal);
+  }
+
+  return DeserializeEncryptedFromStream(properties, input, bloom_filter_length, decryptor,
+                                        row_group_ordinal, column_ordinal);
+}
+
 BlockSplitBloomFilter BlockSplitBloomFilter::Deserialize(
     const ReaderProperties& properties, ArrowInputStream* input,
     std::optional<int64_t> bloom_filter_length) {