GH-34785: [C++][Parquet] Parquet Bloom Filter Writer Implementation

cpp/src/parquet/CMakeLists.txt

@@ -159,6 +159,7 @@ set(PARQUET_SRCS
     arrow/schema_internal.cc
     arrow/writer.cc
     bloom_filter.cc
+    bloom_filter_builder.cc
     bloom_filter_reader.cc
     column_reader.cc
     column_scanner.cc
@@ -367,7 +368,7 @@ install(FILES "${CMAKE_CURRENT_BINARY_DIR}/parquet_version.h"
 add_parquet_test(internals-test
                  SOURCES
                  bloom_filter_test.cc
-                 bloom_filter_reader_test.cc
+                 bloom_filter_reader_writer_test.cc
                  properties_test.cc
                  statistics_test.cc
                  encoding_test.cc

cpp/src/parquet/arrow/arrow_reader_writer_test.cc

@@ -67,6 +67,8 @@
 #include "parquet/arrow/schema.h"
 #include "parquet/arrow/test_util.h"
 #include "parquet/arrow/writer.h"
+#include "parquet/bloom_filter.h"
+#include "parquet/bloom_filter_reader.h"
 #include "parquet/column_writer.h"
 #include "parquet/file_writer.h"
 #include "parquet/page_index.h"
@@ -5648,7 +5650,7 @@ auto encode_double = [](double value) {
 
 }  // namespace
 
-class ParquetPageIndexRoundTripTest : public ::testing::Test {
+class TestingWithPageIndex {
  public:
   void WriteFile(const std::shared_ptr<WriterProperties>& writer_properties,
                  const std::shared_ptr<::arrow::Table>& table) {
@@ -5672,6 +5674,13 @@ class ParquetPageIndexRoundTripTest : public ::testing::Test {
     ASSERT_OK_AND_ASSIGN(buffer_, sink->Finish());
   }
 
+ protected:
+  std::shared_ptr<Buffer> buffer_;
+};
+
+class ParquetPageIndexRoundTripTest : public ::testing::Test,
+                                      public TestingWithPageIndex {
+ public:
   void ReadPageIndexes(int expect_num_row_groups, int expect_num_pages,
                        const std::set<int>& expect_columns_without_index = {}) {
     auto read_properties = default_arrow_reader_properties();
@@ -5740,7 +5749,6 @@ class ParquetPageIndexRoundTripTest : public ::testing::Test {
   }
 
  protected:
-  std::shared_ptr<Buffer> buffer_;
   std::vector<ColumnIndexObject> column_indexes_;
 };
 
@@ -5976,5 +5984,244 @@ TEST_F(ParquetPageIndexRoundTripTest, EnablePerColumn) {
                             /*null_counts=*/{0}}));
 }
 
+class ParquetBloomFilterRoundTripTest : public ::testing::Test,
+                                        public TestingWithPageIndex {
+ public:
+  void ReadBloomFilters(int expect_num_row_groups,
+                        const std::set<int>& expect_columns_without_filter = {}) {
+    auto reader = ParquetFileReader::Open(std::make_shared<BufferReader>(buffer_));
+
+    auto metadata = reader->metadata();
+    ASSERT_EQ(expect_num_row_groups, metadata->num_row_groups());
+
+    auto& bloom_filter_reader = reader->GetBloomFilterReader();
+
+    for (int rg = 0; rg < metadata->num_row_groups(); ++rg) {
+      auto row_group_reader = bloom_filter_reader.RowGroup(rg);
+      ASSERT_NE(row_group_reader, nullptr);
+
+      for (int col = 0; col < metadata->num_columns(); ++col) {
+        bool expect_no_bloom_filter = expect_columns_without_filter.find(col) !=
+                                      expect_columns_without_filter.cend();
+
+        auto bloom_filter = row_group_reader->GetColumnBloomFilter(col);
+        if (expect_no_bloom_filter) {
+          ASSERT_EQ(nullptr, bloom_filter);
+        } else {
+          ASSERT_NE(nullptr, bloom_filter);
+          bloom_filters_.push_back(std::move(bloom_filter));
+        }
+      }
+    }
+  }
+
+  template <typename ArrowType>
+  void VerifyBloomFilterContains(const BloomFilter* bloom_filter,
+                                 const ::arrow::ChunkedArray& chunked_array) {
+    for (auto value : ::arrow::stl::Iterate<ArrowType>(chunked_array)) {
+      if (value == std::nullopt) {
+        continue;
+      }
+      EXPECT_TRUE(bloom_filter->FindHash(bloom_filter->Hash(value.value())));
+    }
+  }
+
+  template <typename ArrowType>
+  void VerifyBloomFilterNotContains(const BloomFilter* bloom_filter,
+                                    const ::arrow::ChunkedArray& chunked_array) {
+    for (auto value : ::arrow::stl::Iterate<ArrowType>(chunked_array)) {
+      if (value == std::nullopt) {
+        continue;
+      }
+      EXPECT_FALSE(bloom_filter->FindHash(bloom_filter->Hash(value.value())));
+    }
+  }
+
+ protected:
+  std::vector<std::unique_ptr<BloomFilter>> bloom_filters_;
+};
+
+TEST_F(ParquetBloomFilterRoundTripTest, SimpleRoundTrip) {
+  auto schema = ::arrow::schema(
+      {::arrow::field("c0", ::arrow::int64()), ::arrow::field("c1", ::arrow::utf8())});
+  BloomFilterOptions options;
+  options.ndv = 10;
+  auto writer_properties = WriterProperties::Builder()
+                               .enable_bloom_filter_options(options, "c0")
+                               ->enable_bloom_filter_options(options, "c1")
+                               ->max_row_group_length(4)
+                               ->build();
+  auto table = ::arrow::TableFromJSON(schema, {R"([
+        [1,     "a"],
+        [2,     "b"],
+        [3,     "c"],
+        [null,  "d"],
+        [5,     null],
+        [6,     "f"]
+  ])"});
+  WriteFile(writer_properties, table);
+
+  ReadBloomFilters(/*expect_num_row_groups=*/2);
+  ASSERT_EQ(4, bloom_filters_.size());
+  std::vector<int64_t> row_group_row_count{4, 2};
+  int64_t current_row = 0;
+  int64_t bloom_filter_idx = 0;  // current index in `bloom_filters_`
+  for (int64_t row_group_id = 0; row_group_id < 2; ++row_group_id) {
+    {
+      // The bloom filter for same column in another row-group.
+      int64_t bloom_filter_idx_another_rg =
+          row_group_id == 0 ? bloom_filter_idx + 2 : bloom_filter_idx - 2;
+      ASSERT_NE(nullptr, bloom_filters_[bloom_filter_idx]);
+      auto col = table->column(0)->Slice(current_row, row_group_row_count[row_group_id]);
+      VerifyBloomFilterContains<::arrow::Int64Type>(
+          bloom_filters_[bloom_filter_idx].get(), *col);
+      VerifyBloomFilterNotContains<::arrow::Int64Type>(
+          bloom_filters_[bloom_filter_idx_another_rg].get(), *col);
+      ++bloom_filter_idx;
+    }
+    {
+      int64_t bloom_filter_idx_another_rg =
+          row_group_id == 0 ? bloom_filter_idx + 2 : bloom_filter_idx - 2;
+      ASSERT_NE(nullptr, bloom_filters_[bloom_filter_idx]);
+      auto col = table->column(1)->Slice(current_row, row_group_row_count[row_group_id]);
+      VerifyBloomFilterContains<::arrow::StringType>(
+          bloom_filters_[bloom_filter_idx].get(), *col);
+      VerifyBloomFilterNotContains<::arrow::StringType>(
+          bloom_filters_[bloom_filter_idx_another_rg].get(), *col);
+      ++bloom_filter_idx;
+    }
+    current_row += row_group_row_count[row_group_id];
+  }
+}
+
+TEST_F(ParquetBloomFilterRoundTripTest, SimpleRoundTripDictionary) {
+  auto origin_schema = ::arrow::schema(
+      {::arrow::field("c0", ::arrow::int64()), ::arrow::field("c1", ::arrow::utf8())});
+  auto schema = ::arrow::schema(
+      {::arrow::field("c0", ::arrow::dictionary(::arrow::int64(), ::arrow::int64())),
+       ::arrow::field("c1", ::arrow::dictionary(::arrow::int64(), ::arrow::utf8()))});
+  bloom_filters_.clear();
+  BloomFilterOptions options;
+  options.ndv = 10;
+  auto writer_properties = WriterProperties::Builder()
+                               .enable_bloom_filter_options(options, "c0")
+                               ->enable_bloom_filter_options(options, "c1")
+                               ->max_row_group_length(4)
+                               ->build();
+  std::vector<std::string> contents = {R"([
+        [1,     "a"],
+        [2,     "b"],
+        [3,     "c"],
+        [null,  "d"],
+        [5,     null],
+        [6,     "f"]
+  ])"};
+  auto dict_encoded_table = ::arrow::TableFromJSON(schema, contents);
+  // using non_dict_table to adapt some interface which doesn't support dictionary.
+  auto table = ::arrow::TableFromJSON(origin_schema, contents);
+  WriteFile(writer_properties, dict_encoded_table);
+
+  ReadBloomFilters(/*expect_num_row_groups=*/2);
+  ASSERT_EQ(4, bloom_filters_.size());
+  std::vector<int64_t> row_group_row_count{4, 2};
+  int64_t current_row = 0;
+  int64_t bloom_filter_idx = 0;  // current index in `bloom_filters_`
+  for (int64_t row_group_id = 0; row_group_id < 2; ++row_group_id) {
+    {
+      // The bloom filter for same column in another row-group.
+      int64_t bloom_filter_idx_another_rg =
+          row_group_id == 0 ? bloom_filter_idx + 2 : bloom_filter_idx - 2;
+      ASSERT_NE(nullptr, bloom_filters_[bloom_filter_idx]);
+      auto col = table->column(0)->Slice(current_row, row_group_row_count[row_group_id]);
+      VerifyBloomFilterContains<::arrow::Int64Type>(
+          bloom_filters_[bloom_filter_idx].get(), *col);
+      VerifyBloomFilterNotContains<::arrow::Int64Type>(
+          bloom_filters_[bloom_filter_idx_another_rg].get(), *col);
+      ++bloom_filter_idx;
+    }
+    {
+      int64_t bloom_filter_idx_another_rg =
+          row_group_id == 0 ? bloom_filter_idx + 2 : bloom_filter_idx - 2;
+      ASSERT_NE(nullptr, bloom_filters_[bloom_filter_idx]);
+      auto col = table->column(1)->Slice(current_row, row_group_row_count[row_group_id]);
+      VerifyBloomFilterContains<::arrow::StringType>(
+          bloom_filters_[bloom_filter_idx].get(), *col);
+      VerifyBloomFilterNotContains<::arrow::StringType>(
+          bloom_filters_[bloom_filter_idx_another_rg].get(), *col);
+      ++bloom_filter_idx;
+    }
+    current_row += row_group_row_count[row_group_id];
+  }
+}
+
+TEST_F(ParquetBloomFilterRoundTripTest, SimpleRoundTripWithOneFilter) {
+  auto schema = ::arrow::schema(
+      {::arrow::field("c0", ::arrow::int64()), ::arrow::field("c1", ::arrow::utf8())});
+  BloomFilterOptions options;
+  options.ndv = 10;
+  auto writer_properties = WriterProperties::Builder()
+                               .enable_bloom_filter_options(options, "c0")
+                               ->disable_bloom_filter("c1")
+                               ->max_row_group_length(4)
+                               ->build();
+  auto table = ::arrow::TableFromJSON(schema, {R"([
+        [1,     "a"],
+        [2,     "b"],
+        [3,     "c"],
+        [null,  "d"],
+        [5,     null],
+        [6,     "f"]
+  ])"});
+  WriteFile(writer_properties, table);
+
+  ReadBloomFilters(/*expect_num_row_groups=*/2, /*expect_columns_without_filter=*/{1});
+  ASSERT_EQ(2, bloom_filters_.size());
+  std::vector<int64_t> row_group_row_count{4, 2};
+  int64_t current_row = 0;
+  int64_t bloom_filter_idx = 0;  // current index in `bloom_filters_`
+  for (int64_t row_group_id = 0; row_group_id < 2; ++row_group_id) {
+    {
+      ASSERT_NE(nullptr, bloom_filters_[bloom_filter_idx]);
+      auto col = table->column(0)->Slice(current_row, row_group_row_count[row_group_id]);
+      VerifyBloomFilterContains<::arrow::Int64Type>(
+          bloom_filters_[bloom_filter_idx].get(), *col);
+      ++bloom_filter_idx;
+    }
+    current_row += row_group_row_count[row_group_id];
+  }
+}
+
+TEST_F(ParquetBloomFilterRoundTripTest, ThrowForBoolean) {
+  auto schema = ::arrow::schema({::arrow::field("boolean_col", ::arrow::boolean())});
+  BloomFilterOptions options;
+  options.ndv = 10;
+  auto writer_properties = WriterProperties::Builder()
+                               .enable_bloom_filter_options(options, "boolean_col")
+                               ->max_row_group_length(4)
+                               ->build();
+  auto table = ::arrow::TableFromJSON(schema, {R"([
+        [true],
+        [null],
+        [false]
+  ])"});
+  std::shared_ptr<SchemaDescriptor> parquet_schema;
+  auto arrow_writer_properties = default_arrow_writer_properties();
+  ASSERT_OK_NO_THROW(ToParquetSchema(schema.get(), *writer_properties,
+                                     *arrow_writer_properties, &parquet_schema));
+  auto schema_node = std::static_pointer_cast<GroupNode>(parquet_schema->schema_root());
+
+  // Write table to buffer.
+  auto sink = CreateOutputStream();
+  auto pool = ::arrow::default_memory_pool();
+  auto writer = ParquetFileWriter::Open(sink, schema_node, writer_properties);
+  std::unique_ptr<FileWriter> arrow_writer;
+  ASSERT_OK(FileWriter::Make(pool, std::move(writer), schema, arrow_writer_properties,
+                             &arrow_writer));
+  auto s = arrow_writer->WriteTable(*table);
+  EXPECT_TRUE(s.IsIOError());
+  EXPECT_THAT(s.message(),
+              ::testing::HasSubstr("BloomFilterBuilder does not support boolean type"));
+}
+
 }  // namespace arrow
 }  // namespace parquet

cpp/src/parquet/bloom_filter.h

@@ -106,6 +106,34 @@ class PARQUET_EXPORT BloomFilter {
   /// @return hash result.
   virtual uint64_t Hash(const FLBA* value, uint32_t len) const = 0;
 
+  // Variant of const reference argument to facilitate template
+
+  /// Compute hash for ByteArray value by using its plain encoding result.
+  ///
+  /// @param value the value to hash.
+  uint64_t Hash(const ByteArray& value) const { return Hash(&value); }
+
+  /// Compute hash for fixed byte array value by using its plain encoding result.
+  ///
+  /// @param value the value to hash.
+  /// @param type_len the value length.
+  uint64_t Hash(const FLBA& value, uint32_t type_len) const {
+    return Hash(&value, type_len);
+  }
+
+  /// Compute hash for Int96 value by using its plain encoding result.
+  ///
+  /// @param value the value to hash.
+  uint64_t Hash(const Int96& value) const { return Hash(&value); }
+
+  /// Compute hash for std::string_view value by using its plain encoding result.
+  ///
+  /// @param value the value to hash.
+  uint64_t Hash(std::string_view value) const {
+    ByteArray ba(value);
+    return Hash(&ba);
+  }
+
   /// Batch compute hashes for 32 bits values by using its plain encoding result.
   ///
   /// @param values values a pointer to the values to hash.