diff --git a/include/paimon/utils/roaring_bitmap32.h b/include/paimon/utils/roaring_bitmap32.h index 41e018a10..0c989a2bb 100644 --- a/include/paimon/utils/roaring_bitmap32.h +++ b/include/paimon/utils/roaring_bitmap32.h @@ -18,6 +18,7 @@ #include #include #include +#include #include #include @@ -135,6 +136,10 @@ class PAIMON_EXPORT RoaringBitmap32 { Iterator End() const; /// @return the iterator moved to the value which is equal or larger than key Iterator EqualOrLarger(int32_t key) const; + /// @return the first value which is equal or larger than x + std::optional NextValue(int32_t x) const; + /// @return the largest value which is smaller than x + std::optional PreviousValue(int32_t x) const; /// Computes the intersection between two bitmaps and returns new bitmap. /// The current bitmap and the provided bitmap are unchanged. diff --git a/src/paimon/common/utils/roaring_bitmap32.cpp b/src/paimon/common/utils/roaring_bitmap32.cpp index de8e90107..1a6a204c9 100644 --- a/src/paimon/common/utils/roaring_bitmap32.cpp +++ b/src/paimon/common/utils/roaring_bitmap32.cpp @@ -310,4 +310,33 @@ RoaringBitmap32::Iterator RoaringBitmap32::EqualOrLarger(int32_t key) const { return iter; } +std::optional RoaringBitmap32::NextValue(int32_t x) const { + auto iter = EqualOrLarger(x); + if (iter == End()) { + return std::nullopt; + } + return *iter; +} + +std::optional RoaringBitmap32::PreviousValue(int32_t x) const { + if (IsEmpty()) { + return std::nullopt; + } + + auto& bitmap = GetRoaringBitmap(roaring_bitmap_); + if (x <= static_cast(bitmap.minimum())) { + return std::nullopt; + } + + const uint64_t rank = bitmap.rank(static_cast(x - 1)); + if (rank == 0) { + return std::nullopt; + } + + uint32_t value = 0; + [[maybe_unused]] bool found = bitmap.select(static_cast(rank - 1), &value); + assert(found); + return static_cast(value); +} + } // namespace paimon diff --git a/src/paimon/common/utils/roaring_bitmap32_test.cpp b/src/paimon/common/utils/roaring_bitmap32_test.cpp index 1abc9ff4c..7276d7355 100644 --- a/src/paimon/common/utils/roaring_bitmap32_test.cpp +++ b/src/paimon/common/utils/roaring_bitmap32_test.cpp @@ -254,6 +254,25 @@ TEST(RoaringBitmap32Test, TestIterator) { ASSERT_EQ(iter, roaring.End()); } +TEST(RoaringBitmap32Test, TestNextAndPreviousValue) { + RoaringBitmap32 roaring = RoaringBitmap32::From({10, 20, 30}); + + ASSERT_EQ(roaring.NextValue(5), std::optional(10)); + ASSERT_EQ(roaring.NextValue(10), std::optional(10)); + ASSERT_EQ(roaring.NextValue(25), std::optional(30)); + ASSERT_EQ(roaring.NextValue(31), std::nullopt); + + ASSERT_EQ(roaring.PreviousValue(10), std::nullopt); + ASSERT_EQ(roaring.PreviousValue(11), std::optional(10)); + ASSERT_EQ(roaring.PreviousValue(20), std::optional(10)); + ASSERT_EQ(roaring.PreviousValue(21), std::optional(20)); + ASSERT_EQ(roaring.PreviousValue(100), std::optional(30)); + + RoaringBitmap32 empty; + ASSERT_EQ(empty.NextValue(0), std::nullopt); + ASSERT_EQ(empty.PreviousValue(0), std::nullopt); +} + TEST(RoaringBitmap32Test, TestIteratorAssignAndMove) { RoaringBitmap32 roaring = RoaringBitmap32::From({10, 100, 200}); diff --git a/src/paimon/format/parquet/file_reader_wrapper.cpp b/src/paimon/format/parquet/file_reader_wrapper.cpp index e7d6bf606..090da9122 100644 --- a/src/paimon/format/parquet/file_reader_wrapper.cpp +++ b/src/paimon/format/parquet/file_reader_wrapper.cpp @@ -164,14 +164,15 @@ void FileReaderWrapper::AdvanceToNextRowGroup() { current_row_group_idx_++; // Skip row groups excluded by read range. while (current_row_group_idx_ < target_row_groups_.size() && - target_row_groups_[current_row_group_idx_].excluded_by_read_range) { + target_row_groups_[current_row_group_idx_].IsExcludedByReadRange()) { current_row_group_idx_++; } if (current_row_group_idx_ >= target_row_groups_.size()) { next_row_to_read_ = num_rows_; } else { next_row_to_read_ = - all_row_group_ranges_[target_row_groups_[current_row_group_idx_].row_group_index].first; + all_row_group_ranges_[target_row_groups_[current_row_group_idx_].GetRowGroupIndex()] + .first; } } @@ -181,10 +182,10 @@ Status FileReaderWrapper::SeekToRow(uint64_t row_number) { filtered_global_offset_ = 0; for (uint64_t i = 0; i < target_row_groups_.size(); i++) { - if (target_row_groups_[i].excluded_by_read_range) { + if (target_row_groups_[i].IsExcludedByReadRange()) { continue; } - int32_t rg_id = target_row_groups_[i].row_group_index; + int32_t rg_id = target_row_groups_[i].GetRowGroupIndex(); uint64_t rg_start = all_row_group_ranges_[rg_id].first; uint64_t rg_end = all_row_group_ranges_[rg_id].second; if (row_number > rg_start && row_number < rg_end) { @@ -200,9 +201,9 @@ Status FileReaderWrapper::SeekToRow(uint64_t row_number) { // Rebuild batch_reader_ for non-page-filtered RGs at/after seek position. std::vector fully_matched_indices; for (uint64_t j = i; j < target_row_groups_.size(); j++) { - if (!target_row_groups_[j].excluded_by_read_range && - !target_row_groups_[j].is_partially_matched) { - fully_matched_indices.push_back(target_row_groups_[j].row_group_index); + if (!target_row_groups_[j].IsExcludedByReadRange() && + !target_row_groups_[j].IsPartiallyMatched()) { + fully_matched_indices.push_back(target_row_groups_[j].GetRowGroupIndex()); } } if (!fully_matched_indices.empty()) { @@ -222,7 +223,7 @@ Status FileReaderWrapper::SeekToRow(uint64_t row_number) { } Result> FileReaderWrapper::NextPageFiltered() { - int32_t rg_id = target_row_groups_[current_row_group_idx_].row_group_index; + int32_t rg_id = target_row_groups_[current_row_group_idx_].GetRowGroupIndex(); // Construct the per-RG streaming reader on demand. if (!current_page_filtered_reader_) { @@ -237,7 +238,7 @@ Result> FileReaderWrapper::NextPageFiltered( file_reader_->parquet_reader(), target_rg, target_column_indices_, page_filtered_read_schema_, file_reader_->properties().cache_options(), pre_buffered, page_ranges, max_chunksize, pool_)); - current_filtered_row_ranges_ = target_rg.row_ranges; + current_filtered_row_ranges_ = target_rg.GetRowRanges(); current_filtered_rg_start_ = all_row_group_ranges_[rg_id].first; filtered_global_offset_ = 0; } @@ -273,7 +274,7 @@ Result> FileReaderWrapper::NextFullyMatched( return std::shared_ptr(); } - int32_t rg_id = target_row_groups_[current_row_group_idx_].row_group_index; + int32_t rg_id = target_row_groups_[current_row_group_idx_].GetRowGroupIndex(); uint64_t rg_end = all_row_group_ranges_[rg_id].second; int64_t num_rows = record_batch->num_rows(); @@ -298,7 +299,7 @@ Result> FileReaderWrapper::Next() { while (current_row_group_idx_ < target_row_groups_.size()) { bool is_partially_matched = - target_row_groups_[current_row_group_idx_].is_partially_matched; + target_row_groups_[current_row_group_idx_].IsPartiallyMatched(); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr batch, is_partially_matched ? NextPageFiltered() : NextFullyMatched()); if (batch) { @@ -368,9 +369,9 @@ std::vector<::arrow::io::ReadRange> FileReaderWrapper::CollectPreBufferRanges( auto file_metadata = file_reader_->parquet_reader()->metadata(); for (const auto& trg : target_row_groups_) { - if (trg.excluded_by_read_range) continue; + if (trg.IsExcludedByReadRange()) continue; - if (trg.is_partially_matched) { + if (trg.IsPartiallyMatched()) { // Page-filtered RGs: only matching page byte ranges. auto page_ranges = PageFilteredRowGroupReader::ComputePageRanges( file_reader_->parquet_reader(), trg, column_indices); @@ -378,7 +379,7 @@ std::vector<::arrow::io::ReadRange> FileReaderWrapper::CollectPreBufferRanges( std::make_move_iterator(page_ranges.end())); } else { // Fully-matched RGs: entire column chunk ranges. - auto rg_metadata = file_metadata->RowGroup(trg.row_group_index); + auto rg_metadata = file_metadata->RowGroup(trg.GetRowGroupIndex()); for (int32_t col_idx : column_indices) { auto col_chunk = rg_metadata->ColumnChunk(col_idx); int64_t offset = col_chunk->data_page_offset(); @@ -416,12 +417,12 @@ Status FileReaderWrapper::PrepareForReading(const std::vector& t std::vector fully_matched_row_groups; uint64_t active_count = 0; for (const auto& trg : target_row_groups_) { - if (trg.excluded_by_read_range) { + if (trg.IsExcludedByReadRange()) { continue; } active_count++; - if (!trg.is_partially_matched) { - fully_matched_row_groups.push_back(trg.row_group_index); + if (!trg.IsPartiallyMatched()) { + fully_matched_row_groups.push_back(trg.GetRowGroupIndex()); } } @@ -455,14 +456,15 @@ Status FileReaderWrapper::PrepareForReading(const std::vector& t // Reset read state. Find the first non-excluded row group. uint64_t first_active_idx = 0; while (first_active_idx < target_row_groups_.size() && - target_row_groups_[first_active_idx].excluded_by_read_range) { + target_row_groups_[first_active_idx].IsExcludedByReadRange()) { first_active_idx++; } if (first_active_idx >= target_row_groups_.size()) { next_row_to_read_ = num_rows_; } else { next_row_to_read_ = - all_row_group_ranges_[target_row_groups_[first_active_idx].row_group_index].first; + all_row_group_ranges_[target_row_groups_[first_active_idx].GetRowGroupIndex()] + .first; } previous_first_row_ = std::numeric_limits::max(); current_row_group_idx_ = first_active_idx; @@ -476,7 +478,7 @@ Status FileReaderWrapper::ApplyReadRanges( const std::vector>& read_ranges) { if (read_ranges.empty()) { for (auto& trg : target_row_groups_) { - trg.excluded_by_read_range = true; + trg.SetExcludedByReadRange(true); } reader_initialized_ = false; return Status::OK(); @@ -492,7 +494,7 @@ Status FileReaderWrapper::ApplyReadRanges( } // Mark each target row group as excluded or not based on the matching set. for (auto& trg : target_row_groups_) { - trg.excluded_by_read_range = matching_rg_indices.count(trg.row_group_index) == 0; + trg.SetExcludedByReadRange(matching_rg_indices.count(trg.GetRowGroupIndex()) == 0); } reader_initialized_ = false; return Status::OK(); diff --git a/src/paimon/format/parquet/file_reader_wrapper.h b/src/paimon/format/parquet/file_reader_wrapper.h index 758ff703a..3221dbf12 100644 --- a/src/paimon/format/parquet/file_reader_wrapper.h +++ b/src/paimon/format/parquet/file_reader_wrapper.h @@ -33,6 +33,7 @@ #include "arrow/type_fwd.h" #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/format/parquet/row_ranges.h" +#include "paimon/format/parquet/target_row_group.h" #include "paimon/result.h" #include "paimon/status.h" #include "parquet/arrow/reader.h" diff --git a/src/paimon/format/parquet/page_filtered_row_group_reader.cpp b/src/paimon/format/parquet/page_filtered_row_group_reader.cpp index 9c87438b8..20c5efb97 100644 --- a/src/paimon/format/parquet/page_filtered_row_group_reader.cpp +++ b/src/paimon/format/parquet/page_filtered_row_group_reader.cpp @@ -234,8 +234,8 @@ Result> PageFilteredRowGroupReader::Re const ::arrow::io::CacheOptions& cache_options, bool pre_buffered, const std::vector<::arrow::io::ReadRange>& page_ranges, int64_t max_chunksize, std::shared_ptr<::arrow::MemoryPool> pool) { - const auto& row_ranges = target_row_group.row_ranges; - int32_t row_group_index = target_row_group.row_group_index; + const auto& row_ranges = target_row_group.GetRowRanges(); + int32_t row_group_index = target_row_group.GetRowGroupIndex(); if (row_ranges.IsEmpty()) { PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr empty_table, @@ -289,8 +289,8 @@ Result> PageFilteredRowGroupReader::Re std::vector<::arrow::io::ReadRange> PageFilteredRowGroupReader::ComputePageRanges( ::parquet::ParquetFileReader* parquet_reader, const TargetRowGroup& target_row_group, const std::vector& column_indices) { - int32_t row_group_index = target_row_group.row_group_index; - const auto& row_ranges = target_row_group.row_ranges; + int32_t row_group_index = target_row_group.GetRowGroupIndex(); + const auto& row_ranges = target_row_group.GetRowRanges(); std::vector<::arrow::io::ReadRange> ranges; auto file_metadata = parquet_reader->metadata(); diff --git a/src/paimon/format/parquet/page_filtered_row_group_reader.h b/src/paimon/format/parquet/page_filtered_row_group_reader.h index 5092bb5ca..c7376512f 100644 --- a/src/paimon/format/parquet/page_filtered_row_group_reader.h +++ b/src/paimon/format/parquet/page_filtered_row_group_reader.h @@ -27,6 +27,7 @@ #include "arrow/record_batch.h" #include "arrow/type.h" #include "paimon/format/parquet/row_ranges.h" +#include "paimon/format/parquet/target_row_group.h" #include "paimon/result.h" #include "parquet/column_reader.h" #include "parquet/file_reader.h" diff --git a/src/paimon/format/parquet/page_filtered_row_group_reader_test.cpp b/src/paimon/format/parquet/page_filtered_row_group_reader_test.cpp index d6bb36ceb..56ba03057 100644 --- a/src/paimon/format/parquet/page_filtered_row_group_reader_test.cpp +++ b/src/paimon/format/parquet/page_filtered_row_group_reader_test.cpp @@ -44,6 +44,7 @@ #include "paimon/status.h" #include "paimon/testing/utils/read_result_collector.h" #include "paimon/testing/utils/testharness.h" +#include "paimon/utils/roaring_bitmap32.h" #include "parquet/arrow/reader.h" #include "parquet/file_reader.h" #include "parquet/properties.h" @@ -129,6 +130,31 @@ class PageFilteredRowGroupReaderTest : public ::testing::Test { paimon::test::ReadResultCollector::CollectResult(batch_reader.get())); } + /// Read back a Parquet file with a predicate, a bitmap, and page index filter enabled. + void ReadWithPredicateAndBitmapImpl(const std::string& file_name, + const std::shared_ptr& read_schema, + const std::shared_ptr& predicate, + const RoaringBitmap32& bitmap, + std::shared_ptr* out, + int32_t batch_size = 1024, + bool enable_page_level_filter = true) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr in, fs_->Open(file_name)); + ASSERT_OK_AND_ASSIGN(int64_t length, in->Length()); + auto in_stream = std::make_shared(in, arrow_pool_, length); + + std::map options; + options[PARQUET_READ_ENABLE_PAGE_INDEX_FILTER] = + enable_page_level_filter ? "true" : "false"; + ASSERT_OK_AND_ASSIGN(auto batch_reader, + ParquetFileBatchReader::Create(std::move(in_stream), options, + batch_size, nullptr, arrow_pool_)); + auto c_schema = std::make_unique(); + ASSERT_TRUE(arrow::ExportSchema(*read_schema, c_schema.get()).ok()); + ASSERT_OK(batch_reader->SetReadSchema(c_schema.get(), predicate, bitmap)); + ASSERT_OK_AND_ASSIGN(*out, + paimon::test::ReadResultCollector::CollectResult(batch_reader.get())); + } + protected: std::shared_ptr arrow_pool_; std::shared_ptr pool_; @@ -939,14 +965,14 @@ TEST_F(PageFilteredRowGroupReaderTest, NestedStructColumnRowGroupFilter) { ASSERT_TRUE(expected->Equals(result->chunk(0))); } -/// Test: Page-level filtering reading the nested struct column along with the predicate column. +/// Test: Page-level filtering reading only the predicate column (no nested column in read schema). /// -/// This verifies that when reading a subset of columns that includes a nested column -/// and the predicate column, the schema mapping and column assembly work correctly. +/// This verifies that when reading only the "id" column (without the nested struct), +/// page-level filtering works correctly since the read schema contains no nested types. /// /// Schema: { id: int32, info: struct } -/// Read schema: { id: int32, info: struct } -/// Predicate on "id": id >= 70. +/// Read schema: { id: int32 } +/// Predicate on "id": id >= 70. Page-level filtering active → rows 70-99 (30 rows). TEST_F(PageFilteredRowGroupReaderTest, NestedStructColumnOnlyReadIdField) { std::string file_name = dir_->Str() + "/nested_struct_only_nested.parquet"; auto data = MakeNestedStructData(100); @@ -1087,6 +1113,95 @@ TEST_F(PageFilteredRowGroupReaderTest, NestedMapColumnRowGroupFilter) { ASSERT_TRUE(expected->Equals(result->chunk(0))); } +/// Test: nested map projection falls back to row-group-level filtering when page index filter is +/// unavailable for nested read schemas. +/// +/// Schema: { id: int32, props: map } +/// 100 rows, 10 per page, 2 row group. +/// Bitmap: {70..99} hits the second row group (50..99). +/// Because nested schema disables page-level filtering, the entire row group 1 (50..99) is read, +/// so rows [50, 99] should all be returned. +TEST_F(PageFilteredRowGroupReaderTest, NestedMapBitmapFallback) { + std::string file_name = dir_->Str() + "/nested_map_projection_fallback.parquet"; + auto data = MakeMapColumnData(100); + WriteTestFile(file_name, data, /*write_batch_size=*/10, /*max_row_group_length=*/50); + + auto field_props = arrow::field("props", arrow::map(arrow::utf8(), arrow::int32())); + auto read_schema = arrow::schema({arrow::field("id", arrow::int32()), field_props}); + + RoaringBitmap32 bitmap; + bitmap.AddRange(70, 100); + + std::shared_ptr result; + ReadWithPredicateAndBitmapImpl(file_name, read_schema, nullptr, bitmap, &result); + + ASSERT_TRUE(result); + // Because page-level filtering is skipped for nested schemas, we read full row groups. + ASSERT_EQ(50, result->length()); + + auto expected = data->Slice(50, 50); + ASSERT_TRUE(expected->Equals(result->chunk(0))); +} + +/// Test: nested list projection falls back to row-group-level filtering when page index filter is +/// unavailable for nested read schemas. +/// +/// Schema: { id: int32, tags: list } +/// 100 rows, 10 per page, 2 row group. +/// Bitmap: {70..99} hits the second row group (50..99). +/// Because nested schema disables page-level filtering, the entire row group 1 (50..99) is read, +/// so rows [50, 99] should all be returned. +TEST_F(PageFilteredRowGroupReaderTest, NestedListBitmapFallback) { + std::string file_name = dir_->Str() + "/nested_list_projection_fallback.parquet"; + auto data = MakeListColumnData(100); + WriteTestFile(file_name, data, /*write_batch_size=*/10, /*max_row_group_length=*/50); + + auto field_tags = arrow::field("tags", arrow::list(arrow::field("item", arrow::int32()))); + auto read_schema = arrow::schema({arrow::field("id", arrow::int32()), field_tags}); + + RoaringBitmap32 bitmap; + bitmap.AddRange(70, 100); + + std::shared_ptr result; + ReadWithPredicateAndBitmapImpl(file_name, read_schema, nullptr, bitmap, &result); + + ASSERT_TRUE(result); + ASSERT_EQ(50, result->length()); + + auto expected = data->Slice(50, 50); + ASSERT_TRUE(expected->Equals(result->chunk(0))); +} + +/// Test: nested struct projection falls back to row-group-level filtering when page index filter is +/// unavailable for nested read schemas. +/// +/// Schema: { id: int32, info: struct } +/// Bitmap: {70..99} hits the second row group (50..99). +/// Because nested schema disables page-level filtering, the entire second row group (50..99) is +/// read. +TEST_F(PageFilteredRowGroupReaderTest, NestedStructBitmapFallback) { + std::string file_name = dir_->Str() + "/nested_struct_projection_fallback.parquet"; + auto data = MakeNestedStructData(100); + WriteTestFile(file_name, data, /*write_batch_size=*/10, /*max_row_group_length=*/50); + + auto field_x = arrow::field("x", arrow::int32()); + auto field_y = arrow::field("y", arrow::int32()); + auto field_info = arrow::field("info", arrow::struct_({field_x, field_y})); + auto read_schema = arrow::schema({arrow::field("id", arrow::int32()), field_info}); + + RoaringBitmap32 bitmap; + bitmap.AddRange(70, 100); + + std::shared_ptr result; + ReadWithPredicateAndBitmapImpl(file_name, read_schema, nullptr, bitmap, &result); + + ASSERT_TRUE(result); + ASSERT_EQ(50, result->length()); + + auto expected = data->Slice(50, 50); + ASSERT_TRUE(expected->Equals(result->chunk(0))); +} + /// Test: rowgroup-level filtering with multiple adjacent nested columns (struct + list). /// /// Schema: { id: int32, info: struct, tags: list } @@ -1143,5 +1258,272 @@ TEST_F(PageFilteredRowGroupReaderTest, MultipleAdjacentNestedColumns) { auto expected = data->Slice(50, 50); ASSERT_TRUE(expected->Equals(result->chunk(0))); } +/// Test: bitmap hits all pages of a subset of row groups (no predicate). +/// +/// 200 rows, 10 rows per page, 100 rows per row group → 2 row groups. +/// RG0: rows 0-99, RG1: rows 100-199. +/// Bitmap: {0..99} hits all pages of RG0, RG1 is excluded entirely. +/// Expected: 100 rows (0-99). +TEST_F(PageFilteredRowGroupReaderTest, BitmapAllPagesSomeRowGroups) { + std::string file_name = dir_->Str() + "/bitmap_all_pages_rg.parquet"; + auto data = MakeSequentialIntData(200); + WriteTestFile(file_name, data, /*write_batch_size=*/10, /*max_row_group_length=*/100); + + RoaringBitmap32 bitmap; + bitmap.AddRange(0, 100); // hits all of RG0 + + auto read_schema = arrow::schema({arrow::field("val", arrow::int32())}); + std::shared_ptr result; + ReadWithPredicateAndBitmapImpl(file_name, read_schema, /*predicate=*/nullptr, bitmap, &result); + ASSERT_TRUE(result); + ASSERT_EQ(100, result->length()); + + auto flat = arrow::Concatenate(result->chunks()).ValueOrDie(); + auto struct_arr = std::dynamic_pointer_cast(flat); + ASSERT_TRUE(struct_arr); + auto val_arr = std::dynamic_pointer_cast(struct_arr->field(0)); + for (int32_t i = 0; i < 100; ++i) { + ASSERT_EQ(i, val_arr->Value(i)); + } +} + +/// Test: bitmap hits partial pages of a row group (no predicate). +/// +/// 200 rows, 10 rows per page, 100 rows per row group → 2 row groups. +/// Bitmap: {30..59} hits pages 3-5 of RG0 (rows 30-59), RG1 excluded. +/// Expected: 30 rows (30-59). +TEST_F(PageFilteredRowGroupReaderTest, BitmapPartialPagesSingleRowGroup) { + std::string file_name = dir_->Str() + "/bitmap_partial_pages_rg.parquet"; + auto data = MakeSequentialIntData(200); + WriteTestFile(file_name, data, /*write_batch_size=*/10, /*max_row_group_length=*/100); + + RoaringBitmap32 bitmap; + bitmap.AddRange(90, 110); // hits pages 3-5 of RG0 + + auto read_schema = arrow::schema({arrow::field("val", arrow::int32())}); + std::shared_ptr result; + ReadWithPredicateAndBitmapImpl(file_name, read_schema, /*predicate=*/nullptr, bitmap, &result); + ASSERT_TRUE(result); + ASSERT_EQ(20, result->length()); + + auto flat = arrow::Concatenate(result->chunks()).ValueOrDie(); + auto struct_arr = std::dynamic_pointer_cast(flat); + ASSERT_TRUE(struct_arr); + auto val_arr = std::dynamic_pointer_cast(struct_arr->field(0)); + for (int32_t i = 0; i < 20; ++i) { + ASSERT_EQ(90 + i, val_arr->Value(i)); + } +} + +/// Test: bitmap hits all pages of some row groups and partial pages of others. +/// +/// 200 rows, 10 rows per page, 100 rows per row group → 2 row groups. +/// Bitmap: {0..99} hits all of RG0 + {120..149} hits pages 2-4 of RG1. +/// Expected: 100 (RG0) + 30 (RG1 partial) = 130 rows. +TEST_F(PageFilteredRowGroupReaderTest, BitmapAllAndPartialPagesMixed) { + std::string file_name = dir_->Str() + "/bitmap_all_and_partial.parquet"; + auto data = MakeSequentialIntData(200); + WriteTestFile(file_name, data, /*write_batch_size=*/10, /*max_row_group_length=*/100); + + RoaringBitmap32 bitmap; + bitmap.AddRange(0, 100); // all of RG0 + bitmap.AddRange(120, 150); // pages 2-4 of RG1 + + auto read_schema = arrow::schema({arrow::field("val", arrow::int32())}); + std::shared_ptr result; + ReadWithPredicateAndBitmapImpl(file_name, read_schema, /*predicate=*/nullptr, bitmap, &result); + ASSERT_TRUE(result); + ASSERT_EQ(130, result->length()); + + // Verify: rows 0-99 + 120-149 + auto flat = arrow::Concatenate(result->chunks()).ValueOrDie(); + auto struct_arr = std::dynamic_pointer_cast(flat); + ASSERT_TRUE(struct_arr); + auto val_arr = std::dynamic_pointer_cast(struct_arr->field(0)); + for (int32_t i = 0; i < 100; ++i) { + ASSERT_EQ(i, val_arr->Value(i)); + } + for (int32_t i = 0; i < 30; ++i) { + ASSERT_EQ(120 + i, val_arr->Value(100 + i)); + } +} + +/// Test: bitmap hits partial pages of a row group, with page-filtered option disabled. +/// +/// 200 rows, 10 rows per page, 100 rows per row group → 2 row groups. +/// Bitmap: {120..149} hits pages 2-4 of RG1. +/// Expected: 100 rows (100-199) because page-filtered option is disabled, so page-level bitmap is +/// ignored. +TEST_F(PageFilteredRowGroupReaderTest, BitmapWithPageFilteredOptionDisabled) { + std::string file_name = dir_->Str() + "/bitmap_all_and_partial.parquet"; + auto data = MakeSequentialIntData(200); + WriteTestFile(file_name, data, /*write_batch_size=*/10, /*max_row_group_length=*/100); + + RoaringBitmap32 bitmap; + bitmap.AddRange(120, 150); // pages 2-4 of RG1 + + auto read_schema = arrow::schema({arrow::field("val", arrow::int32())}); + std::shared_ptr result; + ReadWithPredicateAndBitmapImpl(file_name, read_schema, /*predicate=*/nullptr, bitmap, &result, + 1024, false); + ASSERT_TRUE(result); + ASSERT_EQ(100, result->length()); + + // Verify: 100-199 + auto flat = arrow::Concatenate(result->chunks()).ValueOrDie(); + auto struct_arr = std::dynamic_pointer_cast(flat); + ASSERT_TRUE(struct_arr); + auto val_arr = std::dynamic_pointer_cast(struct_arr->field(0)); + for (int32_t i = 0; i < 100; ++i) { + ASSERT_EQ(100 + i, val_arr->Value(i)); + } +} + +/// Test: bitmap + predicate both applied, bitmap hits all pages of some row groups. +/// +/// 200 rows, 10 rows per page, 100 rows per row group → 2 row groups. +/// Bitmap: {0..99} hits all of RG0. +/// Predicate: val >= 50. Page-level filtering on RG0: pages 5-9. +/// Expected: 50 rows (50-99). +TEST_F(PageFilteredRowGroupReaderTest, BitmapAllPagesWithPredicate) { + std::string file_name = dir_->Str() + "/bitmap_all_predicate.parquet"; + auto data = MakeSequentialIntData(200); + WriteTestFile(file_name, data, /*write_batch_size=*/10, /*max_row_group_length=*/100); + + RoaringBitmap32 bitmap; + bitmap.AddRange(0, 100); // hits all of RG0 + + auto predicate = PredicateBuilder::GreaterOrEqual( + /*field_index=*/0, /*field_name=*/"val", FieldType::INT, Literal(50)); + + auto read_schema = arrow::schema({arrow::field("val", arrow::int32())}); + std::shared_ptr result; + ReadWithPredicateAndBitmapImpl(file_name, read_schema, predicate, bitmap, &result); + ASSERT_TRUE(result); + ASSERT_EQ(50, result->length()); + + auto flat = arrow::Concatenate(result->chunks()).ValueOrDie(); + auto struct_arr = std::dynamic_pointer_cast(flat); + ASSERT_TRUE(struct_arr); + auto val_arr = std::dynamic_pointer_cast(struct_arr->field(0)); + for (int32_t i = 0; i < 50; ++i) { + ASSERT_EQ(50 + i, val_arr->Value(i)); + } +} + +/// Test: bitmap + predicate both applied, bitmap hits partial pages of a row group. +/// +/// 200 rows, 10 rows per page, 100 rows per row group → 2 row groups. +/// Bitmap: {30..59} hits pages 3-5 of RG0 (rows 30-59). +/// Predicate: val >= 40. Page-level filtering further narrows to pages 4-5 (rows 40-59). +/// Expected: 20 rows (40-59). +TEST_F(PageFilteredRowGroupReaderTest, BitmapPartialPagesWithPredicate) { + std::string file_name = dir_->Str() + "/bitmap_partial_predicate.parquet"; + auto data = MakeSequentialIntData(200); + WriteTestFile(file_name, data, /*write_batch_size=*/10, /*max_row_group_length=*/100); + + RoaringBitmap32 bitmap; + bitmap.AddRange(30, 60); // hits pages 3-5 of RG0 + + auto predicate = PredicateBuilder::GreaterOrEqual( + /*field_index=*/0, /*field_name=*/"val", FieldType::INT, Literal(40)); + + auto read_schema = arrow::schema({arrow::field("val", arrow::int32())}); + std::shared_ptr result; + ReadWithPredicateAndBitmapImpl(file_name, read_schema, predicate, bitmap, &result); + ASSERT_TRUE(result); + ASSERT_EQ(20, result->length()); + + auto flat = arrow::Concatenate(result->chunks()).ValueOrDie(); + auto struct_arr = std::dynamic_pointer_cast(flat); + auto val_arr = std::dynamic_pointer_cast(struct_arr->field(0)); + for (int32_t i = 0; i < 20; ++i) { + ASSERT_EQ(40 + i, val_arr->Value(i)); + } +} + +/// Test: bitmap + predicate both applied, bitmap hits all pages of some RG and +/// partial pages of another. +/// +/// 200 rows, 10 rows per page, 100 rows per row group → 2 row groups. +/// Bitmap: {0..99} (all of RG0) + {120..149} (pages 2-4 of RG1). +/// Predicate: val >= 50 AND val < 160. +/// RG0: all pages → page-filtered to val>=50 → rows 50-99 (50 rows) +/// RG1: pages 2-4 (120-149) → page-filtered to val>=50 AND val<160 → all match (30 rows) +/// Expected: 80 rows (50-99 + 120-149). +TEST_F(PageFilteredRowGroupReaderTest, BitmapMixedWithPredicate) { + std::string file_name = dir_->Str() + "/bitmap_mixed_predicate.parquet"; + auto data = MakeSequentialIntData(200); + WriteTestFile(file_name, data, /*write_batch_size=*/10, /*max_row_group_length=*/100); + + RoaringBitmap32 bitmap; + bitmap.AddRange(0, 100); // all of RG0 + bitmap.AddRange(120, 150); // pages 2-4 of RG1 + + ASSERT_OK_AND_ASSIGN( + auto predicate, + PredicateBuilder::And( + {PredicateBuilder::GreaterOrEqual(/*field_index=*/0, /*field_name=*/"val", + FieldType::INT, Literal(50)), + PredicateBuilder::LessThan(/*field_index=*/0, /*field_name=*/"val", FieldType::INT, + Literal(160))})); + + auto read_schema = arrow::schema({arrow::field("val", arrow::int32())}); + std::shared_ptr result; + ReadWithPredicateAndBitmapImpl(file_name, read_schema, predicate, bitmap, &result); + ASSERT_TRUE(result); + ASSERT_EQ(80, result->length()); + + // Verify: rows 50-99 + 120-149 + auto flat = arrow::Concatenate(result->chunks()).ValueOrDie(); + auto struct_arr = std::dynamic_pointer_cast(flat); + ASSERT_TRUE(struct_arr); + auto val_arr = std::dynamic_pointer_cast(struct_arr->field(0)); + for (int32_t i = 0; i < 50; ++i) { + ASSERT_EQ(50 + i, val_arr->Value(i)); + } + for (int32_t i = 0; i < 30; ++i) { + ASSERT_EQ(120 + i, val_arr->Value(50 + i)); + } +} + +/// Test: read parquet with scattered bitmap +/// +/// 200 rows, 50 rows per page, 100 rows per row group → 2 row groups. +/// Bitmap: [20,30), [35, 40), [125, 126), [130, 131), [150, 200) +/// To test if unneeded row at the start and end of pages are filtered out. +/// Expected: 76 rows ([20, 40) + [125, 131) + [150, 200). +TEST_F(PageFilteredRowGroupReaderTest, ScatteredBitmapTest) { + std::string file_name = dir_->Str() + "/scattered_bitmap.parquet"; + auto data = MakeSequentialIntData(200); + WriteTestFile(file_name, data, /*write_batch_size=*/50, /*max_row_group_length=*/100); + + RoaringBitmap32 bitmap; + bitmap.AddRange(20, 30); + bitmap.AddRange(35, 40); + bitmap.Add(125); + bitmap.Add(130); + bitmap.AddRange(150, 200); + + auto read_schema = arrow::schema({arrow::field("val", arrow::int32())}); + std::shared_ptr result; + ReadWithPredicateAndBitmapImpl(file_name, read_schema, /*predicate=*/nullptr, bitmap, &result); + ASSERT_TRUE(result); + ASSERT_EQ(76, result->length()); + + auto flat = arrow::Concatenate(result->chunks()).ValueOrDie(); + auto struct_arr = std::dynamic_pointer_cast(flat); + ASSERT_TRUE(struct_arr); + auto val_arr = std::dynamic_pointer_cast(struct_arr->field(0)); + for (int32_t i = 0; i < 20; ++i) { + ASSERT_EQ(20 + i, val_arr->Value(i)); + } + for (int32_t i = 0; i < 6; ++i) { + ASSERT_EQ(125 + i, val_arr->Value(20 + i)); + } + for (int32_t i = 0; i < 50; ++i) { + ASSERT_EQ(150 + i, val_arr->Value(26 + i)); + } +} } // namespace paimon::parquet::test diff --git a/src/paimon/format/parquet/parquet_file_batch_reader.cpp b/src/paimon/format/parquet/parquet_file_batch_reader.cpp index 43f4c64c8..88ec5abaa 100644 --- a/src/paimon/format/parquet/parquet_file_batch_reader.cpp +++ b/src/paimon/format/parquet/parquet_file_batch_reader.cpp @@ -45,6 +45,7 @@ #include "paimon/core/utils/nested_projection_utils.h" #include "paimon/format/parquet/parquet_field_id_converter.h" #include "paimon/format/parquet/parquet_format_defs.h" +#include "paimon/format/parquet/parquet_schema_util.h" #include "paimon/format/parquet/parquet_timestamp_converter.h" #include "paimon/format/parquet/predicate_converter.h" #include "paimon/reader/batch_reader.h" @@ -157,26 +158,35 @@ Status ParquetFileBatchReader::SetReadSchema( field_index_map[field->name()] = leaf_indices; } - std::vector row_groups = arrow::internal::Iota(reader_->GetNumberOfRowGroups()); + TargetRowGroups target_row_groups = + TargetRowGroup::MakeForAllRowGroups(reader_->GetAllRowGroupRanges()); + PAIMON_ASSIGN_OR_RAISE( + bool enable_page_index_filter, + OptionsUtils::GetValueFromMap(options_, PARQUET_READ_ENABLE_PAGE_INDEX_FILTER, + DEFAULT_PARQUET_READ_ENABLE_PAGE_INDEX_FILTER)); + if (predicate) { - PAIMON_ASSIGN_OR_RAISE(row_groups, - FilterRowGroupsByPredicate(predicate, file_schema, row_groups)); + PAIMON_ASSIGN_OR_RAISE( + target_row_groups, + FilterRowGroupsByPredicate(predicate, file_schema, target_row_groups)); } if (selection_bitmap) { - PAIMON_ASSIGN_OR_RAISE(row_groups, - FilterRowGroupsByBitmap(selection_bitmap.value(), row_groups)); + PAIMON_ASSIGN_OR_RAISE( + target_row_groups, + FilterRowGroupsByBitmap(selection_bitmap.value(), target_row_groups)); + // workaround: page index filter does not support nested fields for now, skip page index + // bitmap pushdown if there is any nested field in the schema + if (!has_nested_field && enable_page_index_filter) { + PAIMON_ASSIGN_OR_RAISE(target_row_groups, + FilterPagesByBitmap(selection_bitmap.value(), + target_row_groups, column_indices)); + } } // Apply page-level filtering after bitmap pruning so we don't read page index // pages for row groups that the bitmap already excluded. - // If no predicate is provided, skip page-level filtering, row_group_row_ranges will be - // empty - std::map row_group_row_ranges; - if (predicate && !row_groups.empty()) { - PAIMON_ASSIGN_OR_RAISE( - bool enable_page_index_filter, - OptionsUtils::GetValueFromMap(options_, PARQUET_READ_ENABLE_PAGE_INDEX_FILTER, - DEFAULT_PARQUET_READ_ENABLE_PAGE_INDEX_FILTER)); - // walkaround: page index filter does not support nested fields for now, skip page index + // If no predicate is provided, skip page-level filtering + if (predicate && !target_row_groups.empty()) { + // workaround: page index filter does not support nested fields for now, skip page index // filter if there is any nested field in the schema if (enable_page_index_filter && !has_nested_field) { // Build column name to index map for page-level filtering. @@ -190,13 +200,9 @@ Status ParquetFileBatchReader::SetReadSchema( column_name_to_index[name] = indices[0]; } } - - std::pair, std::map> page_filter_result; PAIMON_ASSIGN_OR_RAISE( - page_filter_result, - FilterRowGroupsByPageIndex(predicate, column_name_to_index, row_groups)); - row_groups = std::move(page_filter_result.first); - row_group_row_ranges = std::move(page_filter_result.second); + target_row_groups, + FilterRowGroupsByPageIndex(predicate, column_name_to_index, target_row_groups)); } } @@ -204,22 +210,9 @@ Status ParquetFileBatchReader::SetReadSchema( metrics_->SetCounter(ParquetMetrics::READ_ROW_GROUPS_TOTAL, reader_->GetNumberOfRowGroups()); - metrics_->SetCounter(ParquetMetrics::READ_ROW_GROUPS_AFTER_FILTER, row_groups.size()); - - // Build TargetRowGroup list with page-filter info in one shot. - std::vector target_row_groups; - for (int32_t rg_id : row_groups) { - auto it = row_group_row_ranges.find(rg_id); - if (it != row_group_row_ranges.end()) { - target_row_groups.emplace_back(/*rg_index=*/rg_id, /*is_partially_matched=*/true, - /*ranges=*/it->second); - } else { - target_row_groups.emplace_back( - /*rg_index=*/rg_id, /*is_partially_matched=*/false, /*ranges=*/ - RowRanges(Range(0, reader_->GetAllRowGroupRanges()[rg_id].second - - reader_->GetAllRowGroupRanges()[rg_id].first - 1))); - } - } + metrics_->SetCounter(ParquetMetrics::READ_ROW_GROUPS_AFTER_FILTER, + target_row_groups.size()); + PAIMON_RETURN_NOT_OK(UpdateAllTargetRowRanges(target_row_groups)); PAIMON_RETURN_NOT_OK(reader_->PrepareForReadingLazy(target_row_groups, column_indices)); } @@ -227,9 +220,9 @@ Status ParquetFileBatchReader::SetReadSchema( return Status::OK(); } -Result> ParquetFileBatchReader::FilterRowGroupsByPredicate( +Result ParquetFileBatchReader::FilterRowGroupsByPredicate( const std::shared_ptr& predicate, const std::shared_ptr file_schema, - const std::vector& src_row_groups) const { + const TargetRowGroups& src_row_groups) const { if (!predicate) { return Status::Invalid("cannot pushdown an empty predicate"); } @@ -252,58 +245,140 @@ Result> ParquetFileBatchReader::FilterRowGroupsByPredicate( std::shared_ptr file_fragment, parquet_file_format->MakeFragment( file_source, /*partition_expression=*/PredicateConverter::AlwaysTrue(), - /*physical_schema=*/nullptr, /*row_groups=*/src_row_groups)); + /*physical_schema=*/nullptr, + /*row_groups=*/TargetRowGroup::GetRowGroupIndices(src_row_groups))); PAIMON_RETURN_NOT_OK_FROM_ARROW( file_fragment->EnsureCompleteMetadata(reader_->GetFileReader())); PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(arrow::dataset::FragmentVector target_fragments, file_fragment->SplitByRowGroup(bind_expr)); - std::vector target_row_groups; + TargetRowGroups target_row_groups; target_row_groups.reserve(src_row_groups.size()); for (const auto& fragment : target_fragments) { auto parquet_fragment = dynamic_cast(fragment.get()); if (!parquet_fragment) { return Status::Invalid("cannot cast to ParquetFileFragment in ParquetFileBatchReader"); } - target_row_groups.insert(target_row_groups.end(), parquet_fragment->row_groups().begin(), - parquet_fragment->row_groups().end()); + for (auto rg_index : parquet_fragment->row_groups()) { + for (const auto& row_group : src_row_groups) { + if (row_group.GetRowGroupIndex() == rg_index) { + target_row_groups.emplace_back(row_group); + break; + } + } + } } return target_row_groups; } -Result> ParquetFileBatchReader::FilterRowGroupsByBitmap( - const RoaringBitmap32& bitmap, const std::vector& src_row_groups) const { +Result ParquetFileBatchReader::FilterRowGroupsByBitmap( + const RoaringBitmap32& bitmap, const TargetRowGroups& src_row_groups) const { if (bitmap.IsEmpty()) { return Status::Invalid("cannot push down an empty bitmap to ParquetFileBatchReader"); } + const auto& all_row_group_ranges = reader_->GetAllRowGroupRanges(); - // filter row groups by row range - std::vector target_row_groups; - for (const auto& row_group_idx : src_row_groups) { + + TargetRowGroups target_row_groups; + for (const auto& row_group : src_row_groups) { + int32_t row_group_idx = row_group.GetRowGroupIndex(); if (static_cast(row_group_idx) >= all_row_group_ranges.size()) { return Status::Invalid( fmt::format("src row group {} not in row group meta", row_group_idx)); } + // half open interval [start_row_idx, end_row_idx) const auto& [start_row_idx, end_row_idx] = all_row_group_ranges[row_group_idx]; - if (bitmap.ContainsAny(start_row_idx, end_row_idx)) { - target_row_groups.push_back(row_group_idx); + if (!bitmap.ContainsAny(start_row_idx, end_row_idx)) { + continue; } + target_row_groups.emplace_back(row_group); } return target_row_groups; } +Result ParquetFileBatchReader::FilterPagesByBitmap( + const RoaringBitmap32& bitmap, const TargetRowGroups& src_row_groups, + const std::vector& column_indices) const { + auto page_index_reader = reader_->GetPageIndexReader(); + if (!page_index_reader) { + return src_row_groups; + } + + TargetRowGroups target_row_groups; + target_row_groups.reserve(src_row_groups.size()); + for (const auto& row_group : src_row_groups) { + target_row_groups.emplace_back( + FilterRowGroupPagesByBitmap(bitmap, row_group, column_indices, page_index_reader)); + } + return target_row_groups; +} + +TargetRowGroup ParquetFileBatchReader::FilterRowGroupPagesByBitmap( + const RoaringBitmap32& bitmap, const TargetRowGroup& row_group, + const std::vector& column_indices, + const std::shared_ptr<::parquet::PageIndexReader>& page_index_reader) const { + int32_t row_group_idx = row_group.GetRowGroupIndex(); + auto rg_page_index_reader = page_index_reader->RowGroup(row_group_idx); + if (!rg_page_index_reader) { + return row_group; + } + + const auto& all_row_group_ranges = reader_->GetAllRowGroupRanges(); + uint64_t rg_start_row = all_row_group_ranges[row_group_idx].first; + uint64_t rg_row_count = all_row_group_ranges[row_group_idx].second - rg_start_row; + + RowRanges row_ranges = row_group.GetRowRanges(); + for (int32_t col_index : column_indices) { + auto offset_index = rg_page_index_reader->GetOffsetIndex(col_index); + if (!offset_index) { + continue; + } + auto page_ranges = ComputeColumnPageRanges(bitmap, offset_index->page_locations(), + rg_start_row, rg_row_count); + row_ranges = RowRanges::Intersection(row_ranges, page_ranges); + } + if (row_ranges.RowCount() == static_cast(rg_row_count)) { + return row_group; + } else { + return TargetRowGroup(row_group_idx, true, std::move(row_ranges)); + } +} + +RowRanges ParquetFileBatchReader::ComputeColumnPageRanges( + const RoaringBitmap32& bitmap, const std::vector<::parquet::PageLocation>& page_locations, + uint64_t rg_start_row, uint64_t rg_row_count) { + RowRanges page_row_ranges; + for (size_t page_idx = 0; page_idx < page_locations.size(); ++page_idx) { + // half open interval [first_row, last_row) + auto first_row = page_locations[page_idx].first_row_index; + auto last_row = page_idx + 1 < page_locations.size() + ? page_locations[page_idx + 1].first_row_index + : rg_row_count; + + if (!bitmap.ContainsAny(rg_start_row + first_row, rg_start_row + last_row)) { + continue; + } + // closed interval [range_start_row, range_end_row] + auto range_start_row = bitmap.NextValue(rg_start_row + first_row); + auto range_end_row = bitmap.PreviousValue(rg_start_row + last_row); + if (!range_start_row.has_value() || !range_end_row.has_value()) { + continue; + } + page_row_ranges.Add( + Range(range_start_row.value() - rg_start_row, range_end_row.value() - rg_start_row)); + } + return page_row_ranges; +} + // Uses page-level column index statistics to filter row groups and store per-row-group // RowRanges for true page-level skipping. A row group is excluded if ALL its pages are // determined to not match the predicate. For partially matched row groups, RowRanges // are stored for page-level filtering during reading. -Result, std::map>> -ParquetFileBatchReader::FilterRowGroupsByPageIndex( +Result ParquetFileBatchReader::FilterRowGroupsByPageIndex( const std::shared_ptr& predicate, const std::map& column_name_to_index, - const std::vector& src_row_groups) { - std::map rg_row_ranges; - + const TargetRowGroups& src_row_groups) const { if (!predicate) { - return std::make_pair(src_row_groups, rg_row_ranges); + return src_row_groups; } auto page_index_reader = reader_->GetPageIndexReader(); @@ -311,35 +386,41 @@ ParquetFileBatchReader::FilterRowGroupsByPageIndex( PAIMON_LOG_DEBUG(logger_, "Page index not available in file, skipping page-level filtering (%s)", PARQUET_WRITE_ENABLE_PAGE_INDEX); - return std::make_pair(src_row_groups, rg_row_ranges); + return src_row_groups; } auto file_metadata = reader_->GetFileReader()->parquet_reader()->metadata(); - std::vector target_row_groups; - target_row_groups.reserve(src_row_groups.size()); + TargetRowGroups target_row_groups; - for (int32_t row_group_idx : src_row_groups) { + for (const auto& row_group : src_row_groups) { + int32_t row_group_idx = row_group.GetRowGroupIndex(); auto result = reader_->CalculateFilteredRowRanges(row_group_idx, predicate, column_name_to_index); if (!result.ok()) { - target_row_groups.push_back(row_group_idx); + target_row_groups.emplace_back(row_group); continue; } const auto& row_ranges = result.value(); if (!row_ranges.IsEmpty()) { - target_row_groups.push_back(row_group_idx); - int64_t rg_row_count = file_metadata->RowGroup(row_group_idx)->num_rows(); - if (row_ranges.RowCount() < rg_row_count) { - rg_row_ranges[row_group_idx] = row_ranges; + auto intersection = row_group.IsPartiallyMatched() + ? RowRanges::Intersection(row_group.GetRowRanges(), row_ranges) + : row_ranges; + if (intersection.IsEmpty()) { + continue; + } + if (intersection.RowCount() < rg_row_count) { + target_row_groups.emplace_back(row_group_idx, true, intersection); + } else { + target_row_groups.emplace_back(row_group); } } } - return std::make_pair(std::move(target_row_groups), std::move(rg_row_ranges)); + return target_row_groups; } Result ParquetFileBatchReader::NextBatch() { @@ -537,10 +618,10 @@ Status ParquetFileBatchReader::UpdateAllTargetRowRanges( auto all_row_group_ranges = reader_->GetAllRowGroupRanges(); RowRanges all_ranges; for (const auto& target_row_group : target_row_groups) { - for (const auto& range : target_row_group.row_ranges.GetRanges()) { - all_ranges.Add( - Range(range.from + all_row_group_ranges[target_row_group.row_group_index].first, - range.to + all_row_group_ranges[target_row_group.row_group_index].first)); + auto row_group_idx = target_row_group.GetRowGroupIndex(); + for (const auto& range : target_row_group.GetRowRanges().GetRanges()) { + all_ranges.Add(Range(all_row_group_ranges[row_group_idx].first + range.from, + all_row_group_ranges[row_group_idx].first + range.to)); } } all_row_ranges_ = std::move(all_ranges); diff --git a/src/paimon/format/parquet/parquet_file_batch_reader.h b/src/paimon/format/parquet/parquet_file_batch_reader.h index 8a8affce2..e96cffb26 100644 --- a/src/paimon/format/parquet/parquet_file_batch_reader.h +++ b/src/paimon/format/parquet/parquet_file_batch_reader.h @@ -39,6 +39,7 @@ #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/format/parquet/file_reader_wrapper.h" #include "paimon/format/parquet/row_ranges.h" +#include "paimon/format/parquet/target_row_group.h" #include "paimon/logging.h" #include "paimon/reader/prefetch_file_batch_reader.h" #include "paimon/result.h" @@ -191,20 +192,39 @@ class ParquetFileBatchReader : public PrefetchFileBatchReader { Status UpdateAllTargetRowRanges(const std::vector& target_row_groups); // precondition: predicate supposed not be empty - Result> FilterRowGroupsByPredicate( + Result FilterRowGroupsByPredicate( const std::shared_ptr& predicate, const std::shared_ptr file_schema, - const std::vector& src_row_groups) const; - - Result> FilterRowGroupsByBitmap( - const RoaringBitmap32& bitmap, const std::vector& src_row_groups) const; + const TargetRowGroups& src_row_groups) const; + + Result FilterRowGroupsByBitmap(const RoaringBitmap32& bitmap, + const TargetRowGroups& src_row_groups) const; + + Result FilterPagesByBitmap(const RoaringBitmap32& bitmap, + const TargetRowGroups& src_row_groups, + const std::vector& column_indices) const; + + // Apply page-level bitmap filtering to a single row group across all + // requested columns. Intersects the row group's existing ranges with the + // per-column page ranges derived from the bitmap. + TargetRowGroup FilterRowGroupPagesByBitmap( + const RoaringBitmap32& bitmap, const TargetRowGroup& row_group, + const std::vector& column_indices, + const std::shared_ptr<::parquet::PageIndexReader>& page_index_reader) const; + + // Compute the set of row ranges within a single column's pages that + // overlap with the given bitmap. For each page, the bitmap is queried to + // find the first/last matching row in each page, used to trim the page head/tail + static RowRanges ComputeColumnPageRanges( + const RoaringBitmap32& bitmap, const std::vector<::parquet::PageLocation>& page_locations, + uint64_t rg_start_row, uint64_t rg_row_count); // Apply page-level filtering using column index. // Returns (filtered row groups, per-row-group RowRanges for partial matches). - Result, std::map>> - FilterRowGroupsByPageIndex(const std::shared_ptr& predicate, - const std::map& column_name_to_index, - const std::vector& src_row_groups); + Result FilterRowGroupsByPageIndex( + const std::shared_ptr& predicate, + const std::map& column_name_to_index, + const TargetRowGroups& src_row_groups) const; Status GenerateRowMapping(int64_t batch_length); diff --git a/src/paimon/format/parquet/parquet_file_batch_reader_test.cpp b/src/paimon/format/parquet/parquet_file_batch_reader_test.cpp index 3e2f51f58..b18cb8f25 100644 --- a/src/paimon/format/parquet/parquet_file_batch_reader_test.cpp +++ b/src/paimon/format/parquet/parquet_file_batch_reader_test.cpp @@ -195,12 +195,15 @@ class ParquetFileBatchReaderTest : public ::testing::Test, std::unique_ptr PrepareParquetFileBatchReader( const std::string& file_name, const std::shared_ptr& read_schema, const std::shared_ptr& predicate, - const std::optional& selection_bitmap, int32_t batch_size) const { + const std::optional& selection_bitmap, int32_t batch_size, + bool enable_page_level_filter = false) const { EXPECT_OK_AND_ASSIGN(auto input_stream, fs_->Open(file_name)); auto length = fs_->GetFileStatus(file_name).value()->GetLen(); auto in_stream = std::make_unique(std::move(input_stream), pool_, length); - std::map options = {}; + std::map options; + options[PARQUET_READ_ENABLE_PAGE_INDEX_FILTER] = + enable_page_level_filter ? "true" : "false"; return PrepareParquetFileBatchReader(std::move(in_stream), options, read_schema, predicate, selection_bitmap, batch_size); } @@ -715,7 +718,7 @@ TEST_F(ParquetFileBatchReaderTest, TestCreateArrowReaderProperties) { } } -TEST_F(ParquetFileBatchReaderTest, TestBitmapPushDownWithMultiRowGroups) { +TEST_F(ParquetFileBatchReaderTest, TestBitmapRowGroupPushDownWithMultiRowGroups) { arrow::FieldVector fields = {arrow::field("f0", arrow::int32())}; auto arrow_type = arrow::struct_(fields); auto src_array = std::dynamic_pointer_cast( @@ -753,8 +756,47 @@ TEST_F(ParquetFileBatchReaderTest, TestBitmapPushDownWithMultiRowGroups) { auto expected_array = arrow::ChunkedArray::Make({src_array->Slice(0, 6)}).ValueOrDie(); ASSERT_TRUE(result_array->Equals(expected_array)) << result_array->ToString(); } +TEST_F(ParquetFileBatchReaderTest, TestBitmapPagePushDownWithMultiRowGroups) { + arrow::FieldVector fields = {arrow::field("f0", arrow::int32())}; + auto arrow_type = arrow::struct_(fields); + auto src_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(arrow_type, R"([ + [0], + [1], + [2], + [3], + [4], + [5], + [6], + [7], + [8], + [9], + [10], + [11] + ])") + .ValueOrDie()); + auto src_schema = arrow::schema(fields); + std::optional bitmap = RoaringBitmap32::From({3, 5}); + // data in file rowGroup0:[0, 1, 2, 3, 4, 5] | rowGroup1:[6, 7, 8, 9, 10, 11] + + auto arrow_schema = arrow::schema(fields); + WriteArray(file_path_, src_array, arrow_schema, /*write_batch_size=*/12, + /*enable_dictionary=*/true, + /*max_row_group_length=*/6); + + auto parquet_batch_reader = + PrepareParquetFileBatchReader(file_path_, arrow_schema, /*predicate=*/nullptr, bitmap, + /*batch_size=*/12, /*enable_page_level_filter=*/true); + + ASSERT_OK_AND_ASSIGN( + std::shared_ptr result_array, + paimon::test::ReadResultCollector::CollectResult(parquet_batch_reader.get())); + + auto expected_array = arrow::ChunkedArray(src_array->Slice(3, 3)); + ASSERT_TRUE(result_array->Equals(expected_array)) << result_array->ToString(); +} -TEST_F(ParquetFileBatchReaderTest, TestPredicateAndBitmapPushDown) { +TEST_F(ParquetFileBatchReaderTest, TestPredicateAndBitmapRowGroupPushDown) { arrow::FieldVector fields = {arrow::field("f0", arrow::int32())}; auto arrow_type = arrow::struct_(fields); arrow::StructBuilder struct_builder(arrow_type, arrow::default_memory_pool(), @@ -811,6 +853,64 @@ TEST_F(ParquetFileBatchReaderTest, TestPredicateAndBitmapPushDown) { ASSERT_FALSE(result_array); } } +TEST_F(ParquetFileBatchReaderTest, TestPredicateAndBitmapPagePushDown) { + arrow::FieldVector fields = {arrow::field("f0", arrow::int32())}; + auto arrow_type = arrow::struct_(fields); + arrow::StructBuilder struct_builder(arrow_type, arrow::default_memory_pool(), + {std::make_shared()}); + auto int_builder = static_cast(struct_builder.field_builder(0)); + int32_t length = 1024; + for (int32_t i = 0; i < length; ++i) { + ASSERT_TRUE(struct_builder.Append().ok()); + ASSERT_TRUE(int_builder->Append(i).ok()); + } + // data file: + // rowGroup0: [0, 256) + // rowGroup1: [256, 512) + // rowGroup2: [512, 768) + // rowGroup3: [768, 1024) + std::shared_ptr src_array; + ASSERT_TRUE(struct_builder.Finish(&src_array).ok()); + auto src_schema = arrow::schema(fields); + auto arrow_schema = arrow::schema(fields); + WriteArray(file_path_, src_array, arrow_schema, /*write_batch_size=*/1024, + /*enable_dictionary=*/true, + /*max_row_group_length=*/256); + { + // simple case + std::optional bitmap = RoaringBitmap32::From({100, 400, 600}); + ASSERT_OK_AND_ASSIGN( + auto predicate, + PredicateBuilder::Or( + {PredicateBuilder::LessThan(/*field_index=*/0, /*field_name=*/"f0", FieldType::INT, + Literal(255)), + PredicateBuilder::GreaterThan(/*field_index=*/0, /*field_name=*/"f0", + FieldType::INT, Literal(600))})); + auto parquet_batch_reader = + PrepareParquetFileBatchReader(file_path_, arrow_schema, predicate, bitmap, + /*batch_size=*/length, /*enable_page_level_filter=*/true); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr result_array, + paimon::test::ReadResultCollector::CollectResult(parquet_batch_reader.get())); + + auto expected_array = + arrow::ChunkedArray::Make({src_array->Slice(100, 1), src_array->Slice(600, 1)}) + .ValueOrDie(); + ASSERT_TRUE(result_array->Equals(expected_array)) << result_array->ToString(); + } + { + // test all data has been filtered out with predicate and bitmap pushdown + std::optional bitmap = RoaringBitmap32::From({100, 400, 600}); + auto predicate = PredicateBuilder::GreaterThan(/*field_index=*/0, /*field_name=*/"f0", + FieldType::INT, Literal(800)); + auto parquet_batch_reader = PrepareParquetFileBatchReader( + file_path_, arrow_schema, predicate, bitmap, /*batch_size=*/length); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr result_array, + paimon::test::ReadResultCollector::CollectResult(parquet_batch_reader.get())); + ASSERT_FALSE(result_array); + } +} TEST_F(ParquetFileBatchReaderTest, TestReadNoField) { // if only read partition fields, format reader will set empty read schema @@ -1041,7 +1141,8 @@ TEST_F(ParquetFileBatchReaderTest, TestRowMappingSimple) { FieldType::INT, Literal(5), Literal(6))})); auto parquet_batch_reader = PrepareParquetFileBatchReader( - file_path_, arrow_schema, /*predicate=*/predicate, std::nullopt, /*batch_size=*/2); + file_path_, arrow_schema, /*predicate=*/predicate, std::nullopt, /*batch_size=*/2, + /*enable_page_level_filter=*/true); ASSERT_NOK(parquet_batch_reader->GetPreviousBatchFileRowId(0)); ASSERT_OK_AND_ASSIGN( @@ -1107,7 +1208,8 @@ TEST_F(ParquetFileBatchReaderTest, TestRowMappingFullyAndPartially) { FieldType::INT, Literal(8))})); auto parquet_batch_reader = PrepareParquetFileBatchReader( - file_path_, arrow_schema, /*predicate=*/predicate, std::nullopt, /*batch_size=*/3); + file_path_, arrow_schema, /*predicate=*/predicate, std::nullopt, /*batch_size=*/3, + /*enable_page_level_filter=*/true); ASSERT_NOK(parquet_batch_reader->GetPreviousBatchFileRowId(0)); ASSERT_OK_AND_ASSIGN( @@ -1141,7 +1243,8 @@ TEST_F(ParquetFileBatchReaderTest, TestRowMappingSetReadSchemaTwice) { FieldType::INT, Literal(6), Literal(7))})); auto parquet_batch_reader = PrepareParquetFileBatchReader( - file_path_, arrow_schema, /*predicate=*/predicate, std::nullopt, /*batch_size=*/3); + file_path_, arrow_schema, /*predicate=*/predicate, std::nullopt, /*batch_size=*/3, + /*enable_page_level_filter=*/true); ASSERT_NOK(parquet_batch_reader->GetPreviousBatchFileRowId(0)); ASSERT_OK_AND_ASSIGN( diff --git a/src/paimon/format/parquet/row_ranges.h b/src/paimon/format/parquet/row_ranges.h index 6174ac4eb..b2b8338db 100644 --- a/src/paimon/format/parquet/row_ranges.h +++ b/src/paimon/format/parquet/row_ranges.h @@ -105,21 +105,4 @@ class RowRanges { private: std::vector ranges_; }; - -struct TargetRowGroup { - int32_t row_group_index{-1}; - bool is_partially_matched{false}; - - RowRanges row_ranges; - // Whether this row group has been excluded by ApplyReadRanges. - // When true, this row group is logically skipped during iteration - // but retained so that a subsequent wider ApplyReadRanges can restore it. - bool excluded_by_read_range{false}; - - TargetRowGroup() = default; - TargetRowGroup(int32_t rg_index, bool is_partially_matched, RowRanges ranges) - : row_group_index(rg_index), - is_partially_matched(is_partially_matched), - row_ranges(std::move(ranges)) {} -}; } // namespace paimon::parquet diff --git a/src/paimon/format/parquet/target_row_group.h b/src/paimon/format/parquet/target_row_group.h new file mode 100644 index 000000000..3931a6301 --- /dev/null +++ b/src/paimon/format/parquet/target_row_group.h @@ -0,0 +1,112 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include +#include +#include + +#include "paimon/format/parquet/row_ranges.h" + +namespace paimon::parquet { +class TargetRowGroup; +using TargetRowGroups = std::vector; +class TargetRowGroup { + public: + TargetRowGroup(int32_t rg_index, bool is_partially_matched, RowRanges ranges) + : row_group_index_(rg_index), + is_partially_matched_(is_partially_matched), + row_ranges_(std::move(ranges)) {} + + TargetRowGroup(const TargetRowGroup& other) = default; + TargetRowGroup& operator=(const TargetRowGroup& other) = default; + + bool IsExcludedByReadRange() const { + return excluded_by_read_range_; + } + + void SetExcludedByReadRange(bool excluded) { + excluded_by_read_range_ = excluded; + } + + int32_t GetRowGroupIndex() const { + return row_group_index_; + } + + bool IsPartiallyMatched() const { + return is_partially_matched_; + } + + const RowRanges& GetRowRanges() const { + return row_ranges_; + } + + // Create a list of TargetRowGroups for serial (non-filtered) reading. + // + // Each element in 'ranges' is a (start, end) pair describing the row + // range of a single row group. The vector index 'i' is used as the row + // group index, so the caller must ensure that 'ranges' is ordered to + // match the physical row-group order in the file. + // + // For each valid range (start < end), a TargetRowGroup is created with: + // - row_group_index = i + // - is_partially_matched = false (the entire row group is read) + // - row_ranges = [0, end - start - 1] (local row indices covering the + // full group; converted from absolute offsets to 0-based local offsets) + // + // Ranges where start >= end are treated as empty and skipped. + static TargetRowGroups MakeForAllRowGroups( + const std::vector>& ranges) { + TargetRowGroups target_row_groups; + target_row_groups.reserve(ranges.size()); + for (size_t i = 0; i < ranges.size(); ++i) { + // Skip empty or invalid ranges. + if (ranges[i].first >= ranges[i].second) { + continue; + } + // Convert the absolute [start, end) pair into a 0-based local + // row range [0, row_count - 1] for this row group. + target_row_groups.emplace_back( + static_cast(i), false, + RowRanges(Range(0, ranges[i].second - ranges[i].first - 1))); + } + return target_row_groups; + } + + static std::vector GetRowGroupIndices(const TargetRowGroups& target_row_groups) { + std::vector indices; + indices.reserve(target_row_groups.size()); + for (const auto& rg : target_row_groups) { + indices.push_back(rg.GetRowGroupIndex()); + } + return indices; + } + + private: + int32_t row_group_index_{-1}; + bool is_partially_matched_{false}; + // Local row ranges + RowRanges row_ranges_; + // Whether this row group has been excluded by ApplyReadRanges. + // When true, this row group is logically skipped during iteration + // but retained so that a subsequent wider ApplyReadRanges can restore it. + bool excluded_by_read_range_{false}; +}; + +} // namespace paimon::parquet