diff --git a/be/src/format_v2/parquet/parquet_statistics.cpp b/be/src/format_v2/parquet/parquet_statistics.cpp index d0692a80205f3f..41eb8ee553e42d 100644 --- a/be/src/format_v2/parquet/parquet_statistics.cpp +++ b/be/src/format_v2/parquet/parquet_statistics.cpp @@ -397,30 +397,33 @@ bool bloom_filter_excludes(const ParquetColumnSchema& column_schema, int slot_in } struct RowGroupBloomFilterCache { + using CacheKey = std::pair; + ::parquet::BloomFilterReader* bloom_filter_reader = nullptr; - std::map> column_bloom_filters; - std::set loaded_columns; + std::map> column_bloom_filters; + std::set loaded_columns; ::parquet::BloomFilter* get(int row_group_idx, int leaf_column_id, ParquetPruningStats* pruning_stats) { if (bloom_filter_reader == nullptr || leaf_column_id < 0) { return nullptr; } - if (loaded_columns.find(leaf_column_id) == loaded_columns.end()) { - loaded_columns.insert(leaf_column_id); + const CacheKey cache_key {row_group_idx, leaf_column_id}; + if (loaded_columns.find(cache_key) == loaded_columns.end()) { + loaded_columns.insert(cache_key); try { std::shared_ptr<::parquet::RowGroupBloomFilterReader> row_group_reader; if (pruning_stats != nullptr) { SCOPED_RAW_TIMER(&pruning_stats->bloom_filter_read_time); row_group_reader = bloom_filter_reader->RowGroup(row_group_idx); if (row_group_reader != nullptr) { - column_bloom_filters[leaf_column_id] = + column_bloom_filters[cache_key] = row_group_reader->GetColumnBloomFilter(leaf_column_id); } } else { row_group_reader = bloom_filter_reader->RowGroup(row_group_idx); if (row_group_reader != nullptr) { - column_bloom_filters[leaf_column_id] = + column_bloom_filters[cache_key] = row_group_reader->GetColumnBloomFilter(leaf_column_id); } } @@ -430,7 +433,7 @@ struct RowGroupBloomFilterCache { return nullptr; } } - auto it = column_bloom_filters.find(leaf_column_id); + auto it = column_bloom_filters.find(cache_key); return it == column_bloom_filters.end() ? nullptr : it->second.get(); } }; diff --git a/be/test/exec/test_data/parquet_scanner/multi_row_group_bloom_filter.parquet b/be/test/exec/test_data/parquet_scanner/multi_row_group_bloom_filter.parquet new file mode 100644 index 00000000000000..6b9c77842582ad Binary files /dev/null and b/be/test/exec/test_data/parquet_scanner/multi_row_group_bloom_filter.parquet differ diff --git a/be/test/format_v2/parquet/parquet_statistics_test.cpp b/be/test/format_v2/parquet/parquet_statistics_test.cpp index e5cdd8a87457e0..404b50c7c7064e 100644 --- a/be/test/format_v2/parquet/parquet_statistics_test.cpp +++ b/be/test/format_v2/parquet/parquet_statistics_test.cpp @@ -24,6 +24,7 @@ #include #include #include +#include #include #include @@ -270,6 +271,14 @@ VExprContextSPtrs bloom_conjuncts(DataTypePtr data_type, std::vector valu std::make_shared(0, std::move(data_type), std::move(values)))}; } +format::FileScanRequest request_with_bloom_conjunct(DataTypePtr data_type, + std::vector values) { + format::FileScanRequest request; + request.local_positions.emplace(format::LocalColumnId(0), format::LocalIndex(0)); + request.conjuncts = bloom_conjuncts(std::move(data_type), std::move(values)); + return request; +} + void add_bloom_field(segment_v2::BlockSplitBloomFilter* bloom_filter, const Field& value, PrimitiveType type) { DORIS_CHECK(bloom_filter != nullptr); @@ -522,6 +531,41 @@ TEST(ParquetStatisticsPruningTest, VExprUsesDictionaryAndMissingBloomKeepsRows) EXPECT_EQ(pruning_stats.filtered_row_groups_by_bloom_filter, 0); } +TEST(ParquetStatisticsPruningTest, BloomFilterCacheIsScopedToRowGroupAndColumn) { + auto input = arrow::io::ReadableFile::Open( + "./be/test/exec/test_data/parquet_scanner/multi_row_group_bloom_filter.parquet"); + ASSERT_TRUE(input.ok()); + auto reader = ::parquet::ParquetFileReader::Open(*input); + ASSERT_EQ(reader->metadata()->num_row_groups(), 2); + auto& bloom_filter_reader = reader->GetBloomFilterReader(); + for (int row_group_idx = 0; row_group_idx < 2; ++row_group_idx) { + auto row_group_reader = bloom_filter_reader.RowGroup(row_group_idx); + ASSERT_NE(row_group_reader, nullptr); + ASSERT_NE(row_group_reader->GetColumnBloomFilter(0), nullptr); + } + auto schema = build_file_schema(*reader); + + std::vector selected; + format::parquet::ParquetPruningStats pruning_stats; + auto request = request_with_bloom_conjunct(std::make_shared(), + {Field::create_field(12345)}); + ASSERT_TRUE(format::parquet::select_row_groups_by_metadata(*reader->metadata(), reader.get(), + schema, request, nullptr, &selected, + false, &pruning_stats) + .ok()); + EXPECT_EQ(selected, std::vector({0, 1})); + EXPECT_EQ(pruning_stats.filtered_row_groups_by_bloom_filter, 0); + + selected.clear(); + pruning_stats = {}; + ASSERT_TRUE(format::parquet::select_row_groups_by_metadata(*reader->metadata(), reader.get(), + schema, request, nullptr, &selected, + true, &pruning_stats) + .ok()); + EXPECT_EQ(selected, std::vector({1})); + EXPECT_EQ(pruning_stats.filtered_row_groups_by_bloom_filter, 1); +} + TEST(ParquetBloomFilterPruningTest, VExprEqPrunesAbsentIntValue) { auto data_type = std::make_shared(); auto bloom_filter = bloom_filter_for_fields(