This is an automated email from the ASF dual-hosted git repository. SteNicholas pushed a commit to branch release-0.3 in repository https://gitbox.apache.org/repos/asf/paimon-cpp.git
commit c0b8780175a3efed427c5617fd0e7e26e2987c22 Author: Zhou Hongfeng <[email protected]> AuthorDate: Tue Aug 4 18:29:30 2026 +0800 fix(parquet): reuse RowGroupPageIndexReader in FileReaderWrapper layer (#166) --- src/paimon/format/parquet/column_index_filter.cpp | 12 ++--- src/paimon/format/parquet/column_index_filter.h | 8 ++- .../format/parquet/column_index_filter_test.cpp | 49 +++++++++--------- src/paimon/format/parquet/file_reader_wrapper.cpp | 41 +++++++++++---- src/paimon/format/parquet/file_reader_wrapper.h | 9 ++++ .../parquet/page_filtered_row_group_reader.cpp | 24 +++------ .../parquet/page_filtered_row_group_reader.h | 5 +- .../page_filtered_row_group_reader_test.cpp | 58 ++++++++++++++++------ .../format/parquet/parquet_file_batch_reader.cpp | 8 ++- .../format/parquet/parquet_file_batch_reader.h | 7 ++- 10 files changed, 129 insertions(+), 92 deletions(-) diff --git a/src/paimon/format/parquet/column_index_filter.cpp b/src/paimon/format/parquet/column_index_filter.cpp index 6184980..1e0a8f5 100644 --- a/src/paimon/format/parquet/column_index_filter.cpp +++ b/src/paimon/format/parquet/column_index_filter.cpp @@ -38,15 +38,9 @@ namespace paimon::parquet { Result<RowRanges> ColumnIndexFilter::CalculateRowRanges( const std::shared_ptr<Predicate>& predicate, - const std::shared_ptr<::parquet::PageIndexReader>& page_index_reader, - const std::map<std::string, int32_t>& column_name_to_index, int32_t row_group_index, - int64_t row_group_row_count) { - if (!predicate || !page_index_reader) { - return RowRanges::CreateSingle(row_group_row_count); - } - - auto rg_page_index_reader = page_index_reader->RowGroup(row_group_index); - if (!rg_page_index_reader) { + const std::shared_ptr<::parquet::RowGroupPageIndexReader>& rg_page_index_reader, + const std::map<std::string, int32_t>& column_name_to_index, int64_t row_group_row_count) { + if (!predicate || !rg_page_index_reader) { return RowRanges::CreateSingle(row_group_row_count); } diff --git a/src/paimon/format/parquet/column_index_filter.h b/src/paimon/format/parquet/column_index_filter.h index 56bb816..a301624 100644 --- a/src/paimon/format/parquet/column_index_filter.h +++ b/src/paimon/format/parquet/column_index_filter.h @@ -57,16 +57,14 @@ class ColumnIndexFilter { /// Calculate row ranges based on predicate and column indices. /// @param predicate The predicate to evaluate. - /// @param page_index_reader The page index reader for the file. + /// @param rg_page_index_reader The page index reader of target row group for the file. /// @param column_name_to_index Map from column name to column index. - /// @param row_group_index The row group index to filter. /// @param row_group_row_count The number of rows in the row group. /// @return RowRanges that may contain matching rows. static Result<RowRanges> CalculateRowRanges( const std::shared_ptr<Predicate>& predicate, - const std::shared_ptr<::parquet::PageIndexReader>& page_index_reader, - const std::map<std::string, int32_t>& column_name_to_index, int32_t row_group_index, - int64_t row_group_row_count); + const std::shared_ptr<::parquet::RowGroupPageIndexReader>& rg_page_index_reader, + const std::map<std::string, int32_t>& column_name_to_index, int64_t row_group_row_count); private: /// Visit a predicate and calculate row ranges. diff --git a/src/paimon/format/parquet/column_index_filter_test.cpp b/src/paimon/format/parquet/column_index_filter_test.cpp index 35e0257..94c65ef 100644 --- a/src/paimon/format/parquet/column_index_filter_test.cpp +++ b/src/paimon/format/parquet/column_index_filter_test.cpp @@ -305,9 +305,8 @@ class ColumnIndexFilterTest : public ::testing::Test { } Result<RowRanges> Filter(const std::shared_ptr<Predicate>& predicate) { - return ColumnIndexFilter::CalculateRowRanges(predicate, page_index_reader_, - column_name_to_index_, /*row_group_index=*/0, - row_group_row_count_); + return ColumnIndexFilter::CalculateRowRanges(predicate, page_index_reader_->RowGroup(0), + column_name_to_index_, row_group_row_count_); } std::shared_ptr<arrow::MemoryPool> arrow_pool_; @@ -553,19 +552,19 @@ TEST_F(ColumnIndexFilterTest, SignedZeroUsesJavaOrderForFloatingPointPages) { auto less_negative_zero = PredicateBuilder::LessThan( /*field_index=*/0, /*field_name=*/"value", field_type, field_type == FieldType::FLOAT ? Literal(-0.0f) : Literal(-0.0)); - ASSERT_OK_AND_ASSIGN( - auto ranges, ColumnIndexFilter::CalculateRowRanges( - less_negative_zero, page_index_reader, {{"value", 0}}, - /*row_group_index=*/0, reader->metadata()->RowGroup(0)->num_rows())); + ASSERT_OK_AND_ASSIGN(auto ranges, + ColumnIndexFilter::CalculateRowRanges( + less_negative_zero, page_index_reader->RowGroup(0), {{"value", 0}}, + reader->metadata()->RowGroup(0)->num_rows())); ASSERT_TRUE(ranges.IsEmpty()) << "field type: " << static_cast<int32_t>(field_type); auto less_positive_zero = PredicateBuilder::LessThan( /*field_index=*/0, /*field_name=*/"value", field_type, field_type == FieldType::FLOAT ? Literal(0.0f) : Literal(0.0)); - ASSERT_OK_AND_ASSIGN( - ranges, ColumnIndexFilter::CalculateRowRanges( - less_positive_zero, page_index_reader, {{"value", 0}}, - /*row_group_index=*/0, reader->metadata()->RowGroup(0)->num_rows())); + ASSERT_OK_AND_ASSIGN(ranges, + ColumnIndexFilter::CalculateRowRanges( + less_positive_zero, page_index_reader->RowGroup(0), {{"value", 0}}, + reader->metadata()->RowGroup(0)->num_rows())); ASSERT_EQ(20, ranges.RowCount()); ASSERT_EQ(1, ranges.GetRanges().size()); ASSERT_EQ(0, ranges.GetRanges()[0].from); @@ -576,8 +575,8 @@ TEST_F(ColumnIndexFilterTest, SignedZeroUsesJavaOrderForFloatingPointPages) { field_type == FieldType::FLOAT ? Literal(-0.0f) : Literal(-0.0)); ASSERT_OK_AND_ASSIGN( ranges, ColumnIndexFilter::CalculateRowRanges( - greater_negative_zero, page_index_reader, {{"value", 0}}, - /*row_group_index=*/0, reader->metadata()->RowGroup(0)->num_rows())); + greater_negative_zero, page_index_reader->RowGroup(0), {{"value", 0}}, + reader->metadata()->RowGroup(0)->num_rows())); ASSERT_EQ(30, ranges.RowCount()); auto not_equal_negative_zero = PredicateBuilder::NotEqual( @@ -585,17 +584,17 @@ TEST_F(ColumnIndexFilterTest, SignedZeroUsesJavaOrderForFloatingPointPages) { field_type == FieldType::FLOAT ? Literal(-0.0f) : Literal(-0.0)); ASSERT_OK_AND_ASSIGN( ranges, ColumnIndexFilter::CalculateRowRanges( - not_equal_negative_zero, page_index_reader, {{"value", 0}}, - /*row_group_index=*/0, reader->metadata()->RowGroup(0)->num_rows())); + not_equal_negative_zero, page_index_reader->RowGroup(0), {{"value", 0}}, + reader->metadata()->RowGroup(0)->num_rows())); ASSERT_EQ(30, ranges.RowCount()); auto greater_finite = PredicateBuilder::GreaterThan( /*field_index=*/0, /*field_name=*/"value", field_type, field_type == FieldType::FLOAT ? Literal(2.0f) : Literal(2.0)); - ASSERT_OK_AND_ASSIGN( - ranges, ColumnIndexFilter::CalculateRowRanges( - greater_finite, page_index_reader, {{"value", 0}}, - /*row_group_index=*/0, reader->metadata()->RowGroup(0)->num_rows())); + ASSERT_OK_AND_ASSIGN(ranges, + ColumnIndexFilter::CalculateRowRanges( + greater_finite, page_index_reader->RowGroup(0), {{"value", 0}}, + reader->metadata()->RowGroup(0)->num_rows())); ASSERT_TRUE(ranges.IsEmpty()); auto greater_between_pages = PredicateBuilder::GreaterThan( @@ -603,8 +602,8 @@ TEST_F(ColumnIndexFilterTest, SignedZeroUsesJavaOrderForFloatingPointPages) { field_type == FieldType::FLOAT ? Literal(0.5f) : Literal(0.5)); ASSERT_OK_AND_ASSIGN( ranges, ColumnIndexFilter::CalculateRowRanges( - greater_between_pages, page_index_reader, {{"value", 0}}, - /*row_group_index=*/0, reader->metadata()->RowGroup(0)->num_rows())); + greater_between_pages, page_index_reader->RowGroup(0), {{"value", 0}}, + reader->metadata()->RowGroup(0)->num_rows())); ASSERT_EQ(10, ranges.RowCount()); ASSERT_EQ(1, ranges.GetRanges().size()); ASSERT_EQ(20, ranges.GetRanges()[0].from); @@ -613,10 +612,10 @@ TEST_F(ColumnIndexFilterTest, SignedZeroUsesJavaOrderForFloatingPointPages) { auto equal_finite = PredicateBuilder::Equal( /*field_index=*/0, /*field_name=*/"value", field_type, field_type == FieldType::FLOAT ? Literal(2.0f) : Literal(2.0)); - ASSERT_OK_AND_ASSIGN( - ranges, ColumnIndexFilter::CalculateRowRanges( - equal_finite, page_index_reader, {{"value", 0}}, - /*row_group_index=*/0, reader->metadata()->RowGroup(0)->num_rows())); + ASSERT_OK_AND_ASSIGN(ranges, + ColumnIndexFilter::CalculateRowRanges( + equal_finite, page_index_reader->RowGroup(0), {{"value", 0}}, + reader->metadata()->RowGroup(0)->num_rows())); ASSERT_TRUE(ranges.IsEmpty()); } } diff --git a/src/paimon/format/parquet/file_reader_wrapper.cpp b/src/paimon/format/parquet/file_reader_wrapper.cpp index 3966465..4c90b95 100644 --- a/src/paimon/format/parquet/file_reader_wrapper.cpp +++ b/src/paimon/format/parquet/file_reader_wrapper.cpp @@ -232,15 +232,18 @@ Result<std::shared_ptr<arrow::RecordBatch>> FileReaderWrapper::NextPageFiltered( // Construct the per-RG streaming reader on demand. if (!current_page_filtered_reader_) { const auto& target_rg = target_row_groups_[current_row_group_idx_]; + auto row_group_page_index_reader = GetRowGroupPageIndexReader(rg_id); auto page_ranges = PageFilteredRowGroupReader::ComputePageRanges( - target_rg, target_column_indices_, file_reader_->parquet_reader()); + target_rg, target_column_indices_, row_group_page_index_reader, + file_reader_->parquet_reader()); bool pre_buffered = !prebuffered_ranges_.empty(); int64_t max_chunksize = batch_size_ > 0 ? batch_size_ : std::numeric_limits<int64_t>::max(); PAIMON_ASSIGN_OR_RAISE( current_page_filtered_reader_, PageFilteredRowGroupReader::ReadFilteredRowGroup( target_rg, target_column_indices_, file_reader_->properties().cache_options(), - pre_buffered, page_ranges, max_chunksize, pool_, file_reader_.get())); + pre_buffered, page_ranges, max_chunksize, row_group_page_index_reader, pool_, + file_reader_.get())); current_filtered_row_ranges_ = target_rg.GetRowRanges(); current_filtered_rg_start_ = all_row_group_ranges_[rg_id].first; filtered_global_offset_ = 0; @@ -298,6 +301,27 @@ Result<std::shared_ptr<arrow::RecordBatch>> FileReaderWrapper::NextFullyMatched( return record_batch; } +std::shared_ptr<::parquet::RowGroupPageIndexReader> FileReaderWrapper::GetRowGroupPageIndexReader( + int32_t row_group_index) { + auto cached = row_group_page_index_readers_.find(row_group_index); + if (cached != row_group_page_index_readers_.end()) { + return cached->second; + } + + std::shared_ptr<::parquet::RowGroupPageIndexReader> row_group_page_index_reader; + auto page_index_reader = GetPageIndexReader(); + if (page_index_reader) { + row_group_page_index_reader = page_index_reader->RowGroup(row_group_index); + } + + // To avoid OOM, limit the number of row group page index readers cached in memory. + constexpr int32_t kMaxRowGroupPageIndexReaders = 1024; + if (row_group_page_index_readers_.size() < kMaxRowGroupPageIndexReaders) { + row_group_page_index_readers_.emplace(row_group_index, row_group_page_index_reader); + } + return row_group_page_index_reader; +} + Result<std::shared_ptr<arrow::RecordBatch>> FileReaderWrapper::Next() { try { if (PAIMON_UNLIKELY(!reader_initialized_)) { @@ -357,8 +381,9 @@ std::vector<::arrow::io::ReadRange> FileReaderWrapper::CollectPreBufferRanges( if (trg.IsPartiallyMatched()) { // Page-filtered RGs: only matching page byte ranges. + auto row_group_page_index_reader = GetRowGroupPageIndexReader(trg.GetRowGroupIndex()); auto page_ranges = PageFilteredRowGroupReader::ComputePageRanges( - trg, column_indices, file_reader_->parquet_reader()); + trg, column_indices, row_group_page_index_reader, file_reader_->parquet_reader()); ranges.insert(ranges.end(), std::make_move_iterator(page_ranges.begin()), std::make_move_iterator(page_ranges.end())); } else { @@ -500,13 +525,9 @@ Result<RowRanges> FileReaderWrapper::CalculateFilteredRowRanges( return RowRanges::CreateSingle(row_count); } - auto page_index_reader = GetPageIndexReader(); - if (!page_index_reader) { - return RowRanges::CreateSingle(row_count); - } - - return ColumnIndexFilter::CalculateRowRanges( - predicate, page_index_reader, column_name_to_index, row_group_index, row_count); + return ColumnIndexFilter::CalculateRowRanges(predicate, + GetRowGroupPageIndexReader(row_group_index), + column_name_to_index, row_count); } PAIMON_PARQUET_CATCH_AND_RETURN_STATUS("FileReaderWrapper::CalculateFilteredRowRanges") } diff --git a/src/paimon/format/parquet/file_reader_wrapper.h b/src/paimon/format/parquet/file_reader_wrapper.h index 5dbd381..78a07c7 100644 --- a/src/paimon/format/parquet/file_reader_wrapper.h +++ b/src/paimon/format/parquet/file_reader_wrapper.h @@ -143,6 +143,10 @@ class FileReaderWrapper { int32_t row_group_index, const std::shared_ptr<Predicate>& predicate, const std::map<std::string, int32_t>& column_name_to_index); + /// Get or create the page index reader for a row group. + std::shared_ptr<::parquet::RowGroupPageIndexReader> GetRowGroupPageIndexReader( + int32_t row_group_index); + private: FileReaderWrapper(std::unique_ptr<::parquet::arrow::FileReader>&& file_reader, const std::vector<std::pair<uint64_t, uint64_t>>& all_row_group_ranges, @@ -196,6 +200,11 @@ class FileReaderWrapper { // Track pre-buffered ranges so we can wait on destruction std::vector<::arrow::io::ReadRange> prebuffered_ranges_; + + // Arrow caches the file-level PageIndexReader, but RowGroup() creates a new reader each time. + // Keep one reader per row group so its page-index buffers are shared by all read stages. + std::map<int32_t, std::shared_ptr<::parquet::RowGroupPageIndexReader>> + row_group_page_index_readers_; }; } // namespace paimon::parquet 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 c8129db..b3402d4 100644 --- a/src/paimon/format/parquet/page_filtered_row_group_reader.cpp +++ b/src/paimon/format/parquet/page_filtered_row_group_reader.cpp @@ -373,6 +373,7 @@ Result<std::unique_ptr<arrow::RecordBatchReader>> PageFilteredRowGroupReader::Re const TargetRowGroup& target_row_group, const std::vector<int32_t>& column_indices, const ::arrow::io::CacheOptions& cache_options, bool pre_buffered, const std::vector<::arrow::io::ReadRange>& page_ranges, int64_t max_chunksize, + const std::shared_ptr<::parquet::RowGroupPageIndexReader>& row_group_page_index_reader, std::shared_ptr<::arrow::MemoryPool> pool, ::parquet::arrow::FileReader* arrow_file_reader) { auto parquet_reader = arrow_file_reader->parquet_reader(); const auto& row_ranges = target_row_group.GetRowRanges(); @@ -388,14 +389,6 @@ Result<std::unique_ptr<arrow::RecordBatchReader>> PageFilteredRowGroupReader::Re auto rg_metadata = parquet_reader->metadata()->RowGroup(row_group_index); int64_t row_group_row_count = rg_metadata->num_rows(); - // reuse RowGroupPageIndexReader for multiple columns in the same row group to avoid redundant - // metadata reads - std::shared_ptr<::parquet::RowGroupPageIndexReader> rg_page_index_reader; - auto page_index_reader = parquet_reader->GetPageIndexReader(); - if (page_index_reader) { - rg_page_index_reader = page_index_reader->RowGroup(row_group_index); - } - const auto& manifest = arrow_file_reader->manifest(); PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( std::vector<int> field_indices, @@ -407,8 +400,8 @@ Result<std::unique_ptr<arrow::RecordBatchReader>> PageFilteredRowGroupReader::Re for (int field_idx : field_indices) { PAIMON_ASSIGN_OR_RAISE( std::shared_ptr<arrow::ChunkedArray> chunked_array, - ReadFilteredField(rg_page_index_reader, row_group_index, field_idx, column_indices, - row_ranges, row_group_row_count, arrow_file_reader)); + ReadFilteredField(row_group_page_index_reader, row_group_index, field_idx, + column_indices, row_ranges, row_group_row_count, arrow_file_reader)); if (chunked_array->length() != expected_rows) { return Status::Invalid( @@ -434,6 +427,7 @@ Result<std::unique_ptr<arrow::RecordBatchReader>> PageFilteredRowGroupReader::Re std::vector<::arrow::io::ReadRange> PageFilteredRowGroupReader::ComputePageRanges( const TargetRowGroup& target_row_group, const std::vector<int32_t>& column_indices, + const std::shared_ptr<::parquet::RowGroupPageIndexReader>& row_group_page_index_reader, ::parquet::ParquetFileReader* parquet_reader) { int32_t row_group_index = target_row_group.GetRowGroupIndex(); const auto& row_ranges = target_row_group.GetRowRanges(); @@ -447,12 +441,6 @@ std::vector<::arrow::io::ReadRange> PageFilteredRowGroupReader::ComputePageRange auto rg_metadata = file_metadata->RowGroup(row_group_index); int64_t row_group_row_count = rg_metadata->num_rows(); - auto page_index_reader = parquet_reader->GetPageIndexReader(); - std::shared_ptr<::parquet::RowGroupPageIndexReader> rg_page_index_reader; - if (page_index_reader) { - rg_page_index_reader = page_index_reader->RowGroup(row_group_index); - } - for (int32_t col_idx : column_indices) { auto col_chunk = rg_metadata->ColumnChunk(col_idx); const int64_t column_chunk_offset = GetColumnChunkOffset(*col_chunk); @@ -460,8 +448,8 @@ std::vector<::arrow::io::ReadRange> PageFilteredRowGroupReader::ComputePageRange // Try to get OffsetIndex for page-level ranges std::shared_ptr<::parquet::OffsetIndex> offset_index; - if (rg_page_index_reader) { - offset_index = rg_page_index_reader->GetOffsetIndex(col_idx); + if (row_group_page_index_reader) { + offset_index = row_group_page_index_reader->GetOffsetIndex(col_idx); } if (!offset_index) { 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 98303e7..683bde7 100644 --- a/src/paimon/format/parquet/page_filtered_row_group_reader.h +++ b/src/paimon/format/parquet/page_filtered_row_group_reader.h @@ -51,18 +51,20 @@ class PageFilteredRowGroupReader { /// Read a row group with page-level filtering. /// @param target_row_group Target row group with index and row ranges /// @param column_indices Leaf column indices to read - /// @param pool Memory pool /// @param cache_options Cache options for PreBuffer /// @param pre_buffered If true, assumes PreBuffer was already called externally /// and only waits via WhenBuffered (no redundant PreBuffer). /// @param page_ranges If non-empty, wait via WhenBufferedRanges instead of WhenBuffered /// @param max_chunksize Per-batch row cap for the returned reader. + /// @param row_group_page_index_reader Reusable page-index reader for the target row group + /// @param pool Memory pool /// @param arrow_file_reader The Arrow FileReader for ColumnReader tree creation /// @return A RecordBatchReader streaming the filtered rows. static Result<std::unique_ptr<arrow::RecordBatchReader>> ReadFilteredRowGroup( const TargetRowGroup& target_row_group, const std::vector<int32_t>& column_indices, const ::arrow::io::CacheOptions& cache_options, bool pre_buffered, const std::vector<::arrow::io::ReadRange>& page_ranges, int64_t max_chunksize, + const std::shared_ptr<::parquet::RowGroupPageIndexReader>& row_group_page_index_reader, std::shared_ptr<::arrow::MemoryPool> pool, ::parquet::arrow::FileReader* arrow_file_reader); /// Compute the byte ranges of pages that overlap with the given RowRanges. @@ -71,6 +73,7 @@ class PageFilteredRowGroupReader { /// Falls back to entire column chunk range if OffsetIndex is unavailable. static std::vector<::arrow::io::ReadRange> ComputePageRanges( const TargetRowGroup& target_row_group, const std::vector<int32_t>& column_indices, + const std::shared_ptr<::parquet::RowGroupPageIndexReader>& row_group_page_index_reader, ::parquet::ParquetFileReader* parquet_reader); private: 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 a00f895..67b24ac 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 @@ -690,10 +690,13 @@ TEST_F(PageFilteredRowGroupReaderTest, ComputePageRangesPartialMatch) { // Single page match: rows [50, 59] = page 5 RowRanges row_ranges; row_ranges.Add(RowRanges::Range(50, 59)); - + auto page_index_reader = parquet_reader->GetPageIndexReader(); + ASSERT_TRUE(page_index_reader); + auto rg_page_index_reader = page_index_reader->RowGroup(0); auto ranges = PageFilteredRowGroupReader::ComputePageRanges( TargetRowGroup(/*rg_index=*/0, /*is_partially_matched=*/true, /*ranges=*/row_ranges), - /*column_indices=*/{0}, parquet_reader.get()); + /*column_indices=*/{0}, /*row_group_page_index_reader=*/rg_page_index_reader, + parquet_reader.get()); // Should have exactly 1 range (page 5 of column 0, no dictionary since disabled) ASSERT_EQ(1, ranges.size()); @@ -715,10 +718,13 @@ TEST_F(PageFilteredRowGroupReaderTest, ComputePageRangesAllMatch) { // All rows match RowRanges row_ranges; row_ranges.Add(RowRanges::Range(0, 99)); - + auto page_index_reader = parquet_reader->GetPageIndexReader(); + ASSERT_TRUE(page_index_reader); + auto rg_page_index_reader = page_index_reader->RowGroup(0); auto ranges = PageFilteredRowGroupReader::ComputePageRanges( - TargetRowGroup(/*rg_index=*/0, /*is_partially_matched=*/true, /*ranges=*/row_ranges), {0}, - parquet_reader.get()); + TargetRowGroup(/*rg_index=*/0, /*is_partially_matched=*/true, /*ranges=*/row_ranges), + /*column_indices=*/{0}, + /*row_group_page_index_reader=*/rg_page_index_reader, parquet_reader.get()); // 10 pages, all matching ASSERT_EQ(10, ranges.size()); @@ -740,10 +746,13 @@ TEST_F(PageFilteredRowGroupReaderTest, ComputePageRangesNoMatch) { auto parquet_reader = ::parquet::ParquetFileReader::Open(in_stream); RowRanges row_ranges; // empty - + auto page_index_reader = parquet_reader->GetPageIndexReader(); + ASSERT_TRUE(page_index_reader); + auto rg_page_index_reader = page_index_reader->RowGroup(0); auto ranges = PageFilteredRowGroupReader::ComputePageRanges( - TargetRowGroup(/*rg_index=*/0, /*is_partially_matched=*/true, /*ranges=*/row_ranges), {0}, - parquet_reader.get()); + TargetRowGroup(/*rg_index=*/0, /*is_partially_matched=*/true, /*ranges=*/row_ranges), + /*column_indices=*/{0}, + /*row_group_page_index_reader=*/rg_page_index_reader, parquet_reader.get()); ASSERT_EQ(0, ranges.size()); } @@ -763,9 +772,13 @@ TEST_F(PageFilteredRowGroupReaderTest, ComputePageRangesMultiColumn) { RowRanges row_ranges; row_ranges.Add(RowRanges::Range(50, 59)); + auto page_index_reader = parquet_reader->GetPageIndexReader(); + ASSERT_TRUE(page_index_reader); + auto rg_page_index_reader = page_index_reader->RowGroup(0); auto ranges = PageFilteredRowGroupReader::ComputePageRanges( TargetRowGroup(/*rg_index=*/0, /*is_partially_matched=*/true, /*ranges=*/row_ranges), - {0, 1}, parquet_reader.get()); + /*column_indices=*/{0, 1}, /*row_group_page_index_reader=*/rg_page_index_reader, + parquet_reader.get()); // 1 matching page per column = 2 ranges total ASSERT_EQ(2, ranges.size()); @@ -790,9 +803,13 @@ TEST_F(PageFilteredRowGroupReaderTest, ComputePageRangesMultiplePages) { row_ranges.Add(RowRanges::Range(20, 29)); row_ranges.Add(RowRanges::Range(70, 79)); + auto page_index_reader = parquet_reader->GetPageIndexReader(); + ASSERT_TRUE(page_index_reader); + auto rg_page_index_reader = page_index_reader->RowGroup(0); auto ranges = PageFilteredRowGroupReader::ComputePageRanges( - TargetRowGroup(/*rg_index=*/0, /*is_partially_matched=*/true, /*ranges=*/row_ranges), {0}, - parquet_reader.get()); + TargetRowGroup(/*rg_index=*/0, /*is_partially_matched=*/true, /*ranges=*/row_ranges), + /*column_indices=*/{0}, + /*row_group_page_index_reader=*/rg_page_index_reader, parquet_reader.get()); // 2 matching pages for 1 column ASSERT_EQ(2, ranges.size()); @@ -975,9 +992,13 @@ TEST_F(PageFilteredRowGroupReaderTest, ComputePageRangesWithDictionaryEncoding) RowRanges row_ranges; row_ranges.Add(RowRanges::Range(0, 99)); + auto page_index_reader = parquet_reader->GetPageIndexReader(); + ASSERT_TRUE(page_index_reader); + auto rg_page_index_reader = page_index_reader->RowGroup(0); auto ranges = PageFilteredRowGroupReader::ComputePageRanges( TargetRowGroup(/*rg_index=*/0, /*is_partially_matched=*/false, /*ranges=*/row_ranges), - /*column_indices=*/{0}, parquet_reader.get()); + /*column_indices=*/{0}, /*row_group_page_index_reader=*/rg_page_index_reader, + parquet_reader.get()); ASSERT_FALSE(ranges.empty()); @@ -1419,7 +1440,8 @@ TEST_F(PageFilteredRowGroupReaderTest, DirectOffsetIndexJumpReadsEachLeafDiction auto ranges = PageFilteredRowGroupReader::ComputePageRanges( TargetRowGroup(/*rg_index=*/0, /*is_partially_matched=*/true, /*ranges=*/selected_rows), - /*column_indices=*/{0, 1}, parquet_reader.get()); + /*column_indices=*/{0, 1}, /*row_group_page_index_reader=*/row_group_page_index, + parquet_reader.get()); for (int32_t col_idx = 0; col_idx < 2; ++col_idx) { auto column_chunk = row_group->ColumnChunk(col_idx); @@ -1621,11 +1643,16 @@ TEST_F(PageFilteredRowGroupReaderTest, DictionaryEmptySelectionDoesNotReadPages) ASSERT_GE(column_chunk_offset, 0); ASSERT_GT(column_chunk_end, column_chunk_offset); + auto page_index_reader = arrow_file_reader->parquet_reader()->GetPageIndexReader(); + ASSERT_TRUE(page_index_reader); + auto row_group_page_index = page_index_reader->RowGroup(0); + ASSERT_TRUE(row_group_page_index); RowRanges empty_ranges; TargetRowGroup empty_target(/*rg_index=*/0, /*is_partially_matched=*/true, /*ranges=*/empty_ranges); auto page_ranges = PageFilteredRowGroupReader::ComputePageRanges( - empty_target, /*column_indices=*/{0}, arrow_file_reader->parquet_reader()); + empty_target, /*column_indices=*/{0}, /*row_group_page_index_reader=*/row_group_page_index, + arrow_file_reader->parquet_reader()); ASSERT_TRUE(page_ranges.empty()); tracking_input->ClearReadAtRanges(); @@ -1633,7 +1660,8 @@ TEST_F(PageFilteredRowGroupReaderTest, DictionaryEmptySelectionDoesNotReadPages) std::unique_ptr<arrow::RecordBatchReader> result_reader, PageFilteredRowGroupReader::ReadFilteredRowGroup( empty_target, /*column_indices=*/{0}, arrow::io::CacheOptions::Defaults(), - /*pre_buffered=*/false, /*page_ranges=*/{}, /*max_chunksize=*/1024, arrow_pool_, + /*pre_buffered=*/false, /*page_ranges=*/{}, /*max_chunksize=*/1024, + /*row_group_page_index_reader=*/row_group_page_index, arrow_pool_, arrow_file_reader.get())); std::shared_ptr<arrow::RecordBatch> batch; ASSERT_TRUE(result_reader->ReadNext(&batch).ok()); diff --git a/src/paimon/format/parquet/parquet_file_batch_reader.cpp b/src/paimon/format/parquet/parquet_file_batch_reader.cpp index fcc47c1..ab142b7 100644 --- a/src/paimon/format/parquet/parquet_file_batch_reader.cpp +++ b/src/paimon/format/parquet/parquet_file_batch_reader.cpp @@ -408,8 +408,7 @@ Result<TargetRowGroups> ParquetFileBatchReader::RefineRowRangesByTrimming( TargetRowGroups target_row_groups; target_row_groups.reserve(src_row_groups.size()); for (const auto& row_group : src_row_groups) { - auto filtered = - TrimRowGroupPageRanges(bitmap, row_group, column_indices, page_index_reader); + auto filtered = TrimRowGroupPageRanges(bitmap, row_group, column_indices); if (!filtered.GetRowRanges().IsEmpty()) { target_row_groups.emplace_back(std::move(filtered)); } @@ -419,10 +418,9 @@ Result<TargetRowGroups> ParquetFileBatchReader::RefineRowRangesByTrimming( TargetRowGroup ParquetFileBatchReader::TrimRowGroupPageRanges( const RoaringBitmap32& bitmap, const TargetRowGroup& row_group, - const std::vector<int32_t>& column_indices, - const std::shared_ptr<::parquet::PageIndexReader>& page_index_reader) const { + const std::vector<int32_t>& column_indices) const { int32_t row_group_idx = row_group.GetRowGroupIndex(); - auto rg_page_index_reader = page_index_reader->RowGroup(row_group_idx); + auto rg_page_index_reader = reader_->GetRowGroupPageIndexReader(row_group_idx); if (!rg_page_index_reader) { return row_group; } diff --git a/src/paimon/format/parquet/parquet_file_batch_reader.h b/src/paimon/format/parquet/parquet_file_batch_reader.h index 0cdecbf..7e5e9af 100644 --- a/src/paimon/format/parquet/parquet_file_batch_reader.h +++ b/src/paimon/format/parquet/parquet_file_batch_reader.h @@ -217,10 +217,9 @@ class ParquetFileBatchReader : public PrefetchFileBatchReader { // 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 TrimRowGroupPageRanges( - const RoaringBitmap32& bitmap, const TargetRowGroup& row_group, - const std::vector<int32_t>& column_indices, - const std::shared_ptr<::parquet::PageIndexReader>& page_index_reader) const; + TargetRowGroup TrimRowGroupPageRanges(const RoaringBitmap32& bitmap, + const TargetRowGroup& row_group, + const std::vector<int32_t>& column_indices) const; // Apply bitmap filtering to row ranges by coalescing nearby ranges. Result<TargetRowGroups> RefineRowRangesByCoalescing(
