mapleFU commented on code in PR #37400:
URL: https://github.com/apache/arrow/pull/37400#discussion_r1664302177
##########
cpp/src/parquet/arrow/arrow_reader_writer_test.cc:
##########
@@ -5723,5 +5731,213 @@ TEST_F(ParquetPageIndexRoundTripTest, EnablePerColumn) {
/*null_counts=*/{0}}));
}
+class ParquetBloomFilterRoundTripTest : public ::testing::Test,
+ public ParquetIndexRoundTripTest {
+ 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 VerifyBloomFilter(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())));
+ }
+ }
+
+ 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 = 100;
+ 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) {
+ {
+ ASSERT_NE(nullptr, bloom_filters_[bloom_filter_idx]);
+ auto col = table->column(0)->Slice(current_row,
row_group_row_count[row_group_id]);
+
VerifyBloomFilter<::arrow::Int64Type>(bloom_filters_[bloom_filter_idx].get(),
*col);
+ ++bloom_filter_idx;
+ }
+ {
+ ASSERT_NE(nullptr, bloom_filters_[bloom_filter_idx]);
+ auto col = table->column(1)->Slice(current_row,
row_group_row_count[row_group_id]);
+
VerifyBloomFilter<::arrow::StringType>(bloom_filters_[bloom_filter_idx].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 = 100;
+ 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 table = ::arrow::TableFromJSON(schema, contents);
+ auto non_dict_table = ::arrow::TableFromJSON(origin_schema, contents);
Review Comment:
`non_dict_table` is used to adapt some interface which doesn't support
dictionary. I'll add comments here
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscr...@arrow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
