Merge branch 'main' into pk-mor-read-opt

Author: lxy-9602

src/paimon/core/operation/file_store_scan.cpp

@@ -169,8 +169,9 @@ Result<std::shared_ptr<FileStoreScan::RawPlan>> FileStoreScan::CreatePlan() cons
     metrics_->SetCounter(ScanMetrics::LAST_SCANNED_SNAPSHOT_ID,
                          snapshot.has_value() ? snapshot.value().Id() : int64_t{0});
     metrics_->SetCounter(ScanMetrics::LAST_SCANNED_MANIFESTS, filtered_manifest_file_metas.size());
-    metrics_->SetCounter(ScanMetrics::LAST_SCAN_SKIPPED_TABLE_FILES,
-                         all_data_files - manifest_entries.size());
+    metrics_->SetCounter(
+        ScanMetrics::LAST_SCAN_SKIPPED_TABLE_FILES,
+        std::max(int64_t{0}, all_data_files - static_cast<int64_t>(manifest_entries.size())));
     metrics_->SetCounter(ScanMetrics::LAST_SCAN_RESULTED_TABLE_FILES, manifest_entries.size());
     return std::make_shared<FileStoreScan::RawPlan>(scan_mode_, snapshot,
                                                     std::move(manifest_entries));
@@ -219,7 +220,7 @@ Status FileStoreScan::ReadFileEntries(const std::vector<ManifestFileMeta>& manif
                                       std::vector<ManifestEntry>* manifest_entries) const {
     std::vector<std::future<Result<std::vector<ManifestEntry>>>> futures;
     for (const auto& meta : manifest_metas) {
-        auto read_meta_task = [this, &meta]() -> Result<std::vector<ManifestEntry>> {
+        auto read_meta_task = [this, meta]() -> Result<std::vector<ManifestEntry>> {
             std::vector<ManifestEntry> tmp_entries;
             PAIMON_RETURN_NOT_OK(ReadManifestFileMeta(meta, &tmp_entries));
             return tmp_entries;

src/paimon/format/parquet/file_reader_wrapper.cpp

@@ -55,8 +55,8 @@ class FileReaderWrapper {
     ~FileReaderWrapper();
 
     static Result<std::unique_ptr<FileReaderWrapper>> Create(
-        std::unique_ptr<::parquet::arrow::FileReader>&& reader, ::arrow::MemoryPool* pool,
-        int64_t batch_size);
+        std::unique_ptr<::parquet::arrow::FileReader>&& reader, int64_t batch_size,
+        std::shared_ptr<arrow::MemoryPool> pool);
 
     /// Seek to the specified row number.
     /// @param row_number The row to seek to (must be at a row group boundary).
@@ -87,7 +87,7 @@ class FileReaderWrapper {
     }
 
     /// Get the underlying Parquet file reader.
-    ::parquet::arrow::FileReader* GetFileReader() const {
+    ::parquet::arrow::FileReader* GetFileReader() {
         return file_reader_.get();
     }
 
@@ -109,24 +109,18 @@ class FileReaderWrapper {
 
     /// Prepare for lazy reading of the specified row groups and columns.
     /// Actual reader initialization is deferred until the first Next() call.
-    Status PrepareForReadingLazy(const std::set<int32_t>& row_group_indices,
+    Status PrepareForReadingLazy(const std::vector<TargetRowGroup>& target_row_groups,
                                  const std::vector<int32_t>& column_indices);
 
     /// Prepare for immediate reading of the specified row groups and columns.
     /// Initializes the reader and starts pre-buffering I/O.
-    Status PrepareForReading(const std::set<int32_t>& row_group_indices,
+    Status PrepareForReading(const std::vector<TargetRowGroup>& target_row_groups,
                              const std::vector<int32_t>& column_indices);
 
-    /// Filter row groups by read ranges, returning only those that overlap.
-    Result<std::set<int32_t>> FilterRowGroupsByReadRanges(
-        const std::vector<std::pair<uint64_t, uint64_t>>& read_ranges,
-        const std::vector<int32_t>& src_row_groups) const;
-
-    /// Set per-row-group RowRanges for page-level filtering.
-    /// Only partially matched row groups should have entries.
-    void SetRowGroupRowRanges(const std::map<int32_t, RowRanges>& ranges) {
-        row_group_row_ranges_ = ranges;
-    }
+    /// Apply read ranges to the current target_row_groups_, keeping only those
+    /// whose row-group range is equal to one of the given read ranges.
+    /// Resets reader state so that the next Next() call will re-initialize.
+    Status ApplyReadRanges(const std::vector<std::pair<uint64_t, uint64_t>>& read_ranges);
 
     /// Get the page index reader for the file.
     /// Returns nullptr if page index is not available.
@@ -144,21 +138,38 @@ class FileReaderWrapper {
  private:
     FileReaderWrapper(std::unique_ptr<::parquet::arrow::FileReader>&& file_reader,
                       const std::vector<std::pair<uint64_t, uint64_t>>& all_row_group_ranges,
-                      uint64_t num_rows, ::arrow::MemoryPool* pool, int64_t batch_size);
+                      uint64_t num_rows, int64_t batch_size,
+                      std::shared_ptr<::arrow::MemoryPool> pool);
+
+    /// Wait for all pending PreBuffer operations to complete.
+    void WaitForPendingPreBuffer();
+
+    /// Advance current_row_group_idx_ to the next row group and update next_row_to_read_.
+    void AdvanceToNextRowGroup();
+
+    /// Read next batch from a page-filtered row group. Returns nullptr when the RG is exhausted.
+    Result<std::shared_ptr<arrow::RecordBatch>> NextPageFiltered();
 
-    Result<std::set<int32_t>> ReadRangesToRowGroupIds(
-        const std::vector<std::pair<uint64_t, uint64_t>>& read_ranges) const;
-    Result<int32_t> GetRowGroupId(std::pair<uint64_t, uint64_t> target_range) const;
+    /// Read next batch from the fully-matched batch_reader_. Returns nullptr when exhausted.
+    Result<std::shared_ptr<arrow::RecordBatch>> NextFullyMatched();
+
+    /// Build page_filtered_read_schema_ from the given column indices. No-op if already built.
+    Status BuildPageFilteredSchema(const std::vector<int32_t>& column_indices);
+
+    /// Collect all byte ranges that need pre-buffering (page-filtered + fully-matched).
+    std::vector<::arrow::io::ReadRange> CollectPreBufferRanges(
+        const std::vector<int32_t>& column_indices);
+
+    /// Dispatch a single PreBufferRanges call with merged ranges.
+    void DispatchPreBuffer(std::vector<::arrow::io::ReadRange> ranges);
 
     std::unique_ptr<::parquet::arrow::FileReader> file_reader_;
     std::unique_ptr<arrow::RecordBatchReader> batch_reader_;
 
     std::vector<std::pair<uint64_t, uint64_t>> all_row_group_ranges_;
-    std::set<int32_t> target_row_group_indices_;
-    std::vector<std::pair<uint64_t, uint64_t>> target_row_groups_;
     std::vector<int32_t> target_column_indices_;
 
-    ::arrow::MemoryPool* pool_;
+    std::shared_ptr<::arrow::MemoryPool> pool_;
     int64_t batch_size_;  // 0 means no limit
 
     const uint64_t num_rows_;
@@ -175,13 +186,8 @@ class FileReaderWrapper {
     RowRanges current_filtered_row_ranges_;   // RowRanges for the active page-filtered RG
     uint64_t current_filtered_rg_start_ = 0;  // Absolute row-group start row number
 
-    // Page-level filtering state. Externally injected via SetRowGroupRowRanges and
-    // looked up by row group index when entering a page-filtered RG.
-    std::map<int32_t, RowRanges> row_group_row_ranges_;
-
-    // Set of target_row_groups_ positional indices that use page-filtered reading.
-    // Built in PrepareForReading from row_group_row_ranges_.
-    std::set<uint64_t> page_filtered_indices_;
+    // Target row groups with row ranges for none page-level filtering and page-level filtering
+    std::vector<TargetRowGroup> target_row_groups_;
 
     // Arrow schema covering target_column_indices_, used when constructing the per-RG
     // page-filtered reader. Cached in PrepareForReading because it's identical across
@@ -190,9 +196,6 @@ class FileReaderWrapper {
 
     // Track pre-buffered ranges so we can wait on destruction
     std::vector<::arrow::io::ReadRange> prebuffered_ranges_;
-
-    /// Wait for all pending PreBuffer operations to complete.
-    void WaitForPendingPreBuffer();
 };
 
 }  // namespace paimon::parquet

src/paimon/format/parquet/file_reader_wrapper.h

@@ -135,8 +135,7 @@ class FileReaderWrapperTest : public ::testing::Test {
         PAIMON_RETURN_NOT_OK_FROM_ARROW(file_reader_builder.memory_pool(arrow_pool_.get())
                                             ->properties(arrow_reader_props)
                                             ->Build(&file_reader));
-        return FileReaderWrapper::Create(std::move(file_reader), ::arrow::default_memory_pool(),
-                                         wrapper_batch_size);
+        return FileReaderWrapper::Create(std::move(file_reader), wrapper_batch_size, arrow_pool_);
     }
 
     void PrepareParquetFile(const std::string& file_path, int32_t row_count,
@@ -196,8 +195,9 @@ TEST_F(FileReaderWrapperTest, EmptyFile) {
 }
 
 TEST_F(FileReaderWrapperTest, NullFileReader) {
-    ASSERT_NOK_WITH_MSG(FileReaderWrapper::Create(nullptr, ::arrow::default_memory_pool(),
-                                                  /*batch_size=*/0),
+    ASSERT_NOK_WITH_MSG(FileReaderWrapper::Create(nullptr,
+                                                  /*batch_size=*/0,
+                                                  /*pool=*/arrow_pool_),
                         "file reader wrapper create failed. file reader is nullptr");
 }
 
@@ -261,11 +261,11 @@ TEST_F(FileReaderWrapperTest, PageFilteredZeroBatchSizeDoesNotHang) {
     // contiguous ranges keep the test honest about RowRanges semantics; the actual
     // numbers don't matter as long as their total falls inside the row group.
     RowRanges rr({RowRanges::Range(0, 49), RowRanges::Range(100, 149)});
-    reader_wrapper->SetRowGroupRowRanges({{0, rr}});
 
     std::vector<int32_t> all_columns = {0, 1, 2};
-    ASSERT_OK(reader_wrapper->PrepareForReading({0}, all_columns));
-
+    ASSERT_OK(reader_wrapper->PrepareForReading(
+        {TargetRowGroup(/*rg_index=*/0, /*is_partially_matched=*/true, /*ranges=*/rr)},
+        all_columns));
     int64_t total = 0;
     int64_t batch_count = 0;
     while (true) {
@@ -295,10 +295,14 @@ TEST_F(FileReaderWrapperTest, SeekBackToConsumedPageFilteredRowGroup) {
     std::map<int32_t, RowRanges> row_ranges_map;
     row_ranges_map[0] = RowRanges(RowRanges::Range(10, 49));
     row_ranges_map[1] = RowRanges(RowRanges::Range(100, 149));
-    reader_wrapper->SetRowGroupRowRanges(row_ranges_map);
 
     std::vector<int32_t> all_columns = {0, 1, 2};
-    ASSERT_OK(reader_wrapper->PrepareForReading({0, 1}, all_columns));
+    ASSERT_OK(reader_wrapper->PrepareForReading(
+        {TargetRowGroup(/*rg_index=*/0, /*is_partially_matched=*/true,
+                        /*ranges=*/row_ranges_map[0]),
+         TargetRowGroup(/*rg_index=*/1, /*is_partially_matched=*/true,
+                        /*ranges=*/row_ranges_map[1])},
+        all_columns));
 
     auto count_all_rows = [&](int64_t* out_total) {
         int64_t total = 0;
@@ -348,8 +352,9 @@ TEST_F(FileReaderWrapperTest, PageFilteredRespectsBatchSize) {
     for (int64_t batch_size : {int64_t{1}, int64_t{2}, int64_t{3}, int64_t{5}, int64_t{10}}) {
         SCOPED_TRACE("batch_size=" + std::to_string(batch_size));
         ASSERT_OK_AND_ASSIGN(auto reader_wrapper, PrepareReaderWrapper(file_path, batch_size));
-        reader_wrapper->SetRowGroupRowRanges({{0, rr}});
-        ASSERT_OK(reader_wrapper->PrepareForReading({0}, {0, 1, 2}));
+        ASSERT_OK(reader_wrapper->PrepareForReading(
+            {TargetRowGroup(/*rg_index=*/0, /*is_partially_matched=*/true, /*ranges=*/rr)},
+            {0, 1, 2}));
 
         int64_t total = 0;
         int64_t batch_count = 0;
@@ -380,45 +385,93 @@ TEST_F(FileReaderWrapperTest, GetRowGroupRanges) {
     ASSERT_TRUE(ranges.empty());
 }
 
-TEST_F(FileReaderWrapperTest, ReadRangesToRowGroupIds) {
+TEST_F(FileReaderWrapperTest, ApplyReadRanges) {
     std::string file_path = PathUtil::JoinPath(dir_->Str(), "test.parquet");
     PrepareParquetFile(file_path, /*row_count=*/5500);
     ASSERT_OK_AND_ASSIGN(auto reader_wrapper, PrepareReaderWrapper(file_path));
-    std::set<int32_t> expected_row_group_ids = {0, 3, 5};
+
+    // Prepare with a subset of row groups: {0, 1, 2, 4, 5}
+    std::vector<TargetRowGroup> initial_targets = {
+        TargetRowGroup(/*rg_index=*/0, /*is_partially_matched=*/false,
+                       /*ranges=*/RowRanges()),
+        TargetRowGroup(/*rg_index=*/1, /*is_partially_matched=*/false,
+                       /*ranges=*/RowRanges()),
+        TargetRowGroup(/*rg_index=*/2, /*is_partially_matched=*/false,
+                       /*ranges=*/RowRanges()),
+        TargetRowGroup(/*rg_index=*/4, /*is_partially_matched=*/false,
+                       /*ranges=*/RowRanges()),
+        TargetRowGroup(/*rg_index=*/5, /*is_partially_matched=*/false,
+                       /*ranges=*/RowRanges())};
+    std::vector<int32_t> all_columns = {0, 1, 2};
+    ASSERT_OK(reader_wrapper->PrepareForReadingLazy(initial_targets, all_columns));
+
+    // Apply read ranges that match RG 0, 3, 5. Only 0 and 5 are in initial targets.
     std::vector<std::pair<uint64_t, uint64_t>> read_ranges = {
         {0, 1000}, {3000, 4000}, {5000, 5500}};
-    ASSERT_OK_AND_ASSIGN(auto row_group_ids, reader_wrapper->ReadRangesToRowGroupIds(read_ranges));
-    ASSERT_EQ(expected_row_group_ids, row_group_ids);
-    std::vector<std::pair<uint64_t, uint64_t>> invalid_ranges = {
-        {0, 1000}, {3000, 4000}, {5000, 5600}};
-    ASSERT_NOK_WITH_MSG(reader_wrapper->ReadRangesToRowGroupIds(invalid_ranges),
-                        "not match with row group range bound");
-    ASSERT_OK_AND_ASSIGN(row_group_ids, reader_wrapper->ReadRangesToRowGroupIds({}));
-    ASSERT_TRUE(row_group_ids.empty());
+    ASSERT_OK(reader_wrapper->ApplyReadRanges(read_ranges));
+
+    // Verify: reading should only produce rows from RG 0 (1000 rows) and RG 5 (500 rows).
+    int64_t total_rows = 0;
+    while (true) {
+        ASSERT_OK_AND_ASSIGN(auto batch, reader_wrapper->Next());
+        if (!batch) {
+            break;
+        }
+        total_rows += batch->num_rows();
+    }
+    ASSERT_EQ(1500, total_rows);
+
+    // Apply empty read ranges should result in no data.
+    ASSERT_OK(reader_wrapper->PrepareForReadingLazy(initial_targets, all_columns));
+    ASSERT_OK(reader_wrapper->ApplyReadRanges({}));
+    ASSERT_OK_AND_ASSIGN(auto batch, reader_wrapper->Next());
+    ASSERT_FALSE(batch);
 }
 
-TEST_F(FileReaderWrapperTest, FilterRowGroupsByReadRanges) {
+TEST_F(FileReaderWrapperTest, ApplyReadRangesWiderSecondCall) {
     std::string file_path = PathUtil::JoinPath(dir_->Str(), "test.parquet");
     PrepareParquetFile(file_path, /*row_count=*/5500);
     ASSERT_OK_AND_ASSIGN(auto reader_wrapper, PrepareReaderWrapper(file_path));
-    std::set<int32_t> expected_row_group_ids = {0, 5};
-    std::vector<std::pair<uint64_t, uint64_t>> read_ranges = {
-        {0, 1000}, {3000, 4000}, {5000, 5500}};
-    ASSERT_OK_AND_ASSIGN(auto row_group_ids,
-                         reader_wrapper->FilterRowGroupsByReadRanges(read_ranges, {0, 1, 2, 4, 5}));
-    ASSERT_EQ(expected_row_group_ids, row_group_ids);
 
-    ASSERT_OK_AND_ASSIGN(row_group_ids,
-                         reader_wrapper->FilterRowGroupsByReadRanges(read_ranges, {}));
-    ASSERT_TRUE(row_group_ids.empty());
+    // Prepare with row groups: {0, 1, 2, 4, 5}
+    std::vector<TargetRowGroup> initial_targets = {
+        TargetRowGroup(/*rg_index=*/0, /*is_partially_matched=*/false,
+                       /*ranges=*/RowRanges()),
+        TargetRowGroup(/*rg_index=*/1, /*is_partially_matched=*/false,
+                       /*ranges=*/RowRanges()),
+        TargetRowGroup(/*rg_index=*/2, /*is_partially_matched=*/false,
+                       /*ranges=*/RowRanges()),
+        TargetRowGroup(/*rg_index=*/4, /*is_partially_matched=*/false,
+                       /*ranges=*/RowRanges()),
+        TargetRowGroup(/*rg_index=*/5, /*is_partially_matched=*/false,
+                       /*ranges=*/RowRanges())};
+    std::vector<int32_t> all_columns = {0, 1, 2};
+    ASSERT_OK(reader_wrapper->PrepareForReadingLazy(initial_targets, all_columns));
+
+    // First ApplyReadRanges: narrow to RG 0 only.
+    ASSERT_OK(reader_wrapper->ApplyReadRanges({{0, 1000}}));
+
+    // Second ApplyReadRanges: widen to RG 0, 1, 2. Previously excluded RG 1, 2 should restore.
+    ASSERT_OK(reader_wrapper->ApplyReadRanges({{0, 1000}, {1000, 2000}, {2000, 3000}}));
+
+    // Verify: reading should produce rows from RG 0 + 1 + 2 = 3000 rows.
+    int64_t total_rows = 0;
+    while (true) {
+        ASSERT_OK_AND_ASSIGN(auto batch, reader_wrapper->Next());
+        if (!batch) break;
+        total_rows += batch->num_rows();
+    }
+    ASSERT_EQ(3000, total_rows);
 }
 
 TEST_F(FileReaderWrapperTest, PrepareForReading) {
     std::string file_path = PathUtil::JoinPath(dir_->Str(), "test.parquet");
     PrepareParquetFile(file_path, /*row_count=*/5500);
     ASSERT_OK_AND_ASSIGN(auto reader_wrapper, PrepareReaderWrapper(file_path));
-    ASSERT_OK(reader_wrapper->PrepareForReading(/*row_group_indices=*/{1},
-                                                /*column_indices=*/{0}));
+    ASSERT_OK(reader_wrapper->PrepareForReading(
+        /*target_row_groups=*/{TargetRowGroup(/*rg_index=*/1, /*is_partially_matched=*/false,
+                                              /*ranges=*/RowRanges())},
+        /*column_indices=*/{0}));
     // seek before actual read range
     ASSERT_OK(reader_wrapper->SeekToRow(0));
     ASSERT_EQ(1000, reader_wrapper->GetNextRowToRead());
@@ -438,8 +491,12 @@ TEST_F(FileReaderWrapperTest, PrepareForReading) {
     ASSERT_FALSE(record_batch);
 
     // empty column indices
-    ASSERT_OK(reader_wrapper->PrepareForReading(/*row_group_indices=*/{0, 1},
-                                                /*column_indices=*/{}));
+    ASSERT_OK(reader_wrapper->PrepareForReading(
+        /*target_row_groups=*/{TargetRowGroup(/*rg_index=*/0, /*is_partially_matched=*/false,
+                                              /*ranges=*/RowRanges()),
+                               TargetRowGroup(/*rg_index=*/1, /*is_partially_matched=*/false,
+                                              /*ranges=*/RowRanges())},
+        /*column_indices=*/{}));
     ASSERT_EQ(0, reader_wrapper->GetNextRowToRead());
     ASSERT_EQ(std::numeric_limits<uint64_t>::max(),
               reader_wrapper->GetPreviousBatchFirstRowNumber().value());
@@ -448,8 +505,9 @@ TEST_F(FileReaderWrapperTest, PrepareForReading) {
     ASSERT_EQ(0, record_batch->num_columns());
 
     // empty row group indices
-    ASSERT_OK(reader_wrapper->PrepareForReading(/*row_group_indices=*/{},
-                                                /*column_indices=*/{0}));
+    ASSERT_OK(reader_wrapper->PrepareForReading(
+        /*target_row_groups=*/{},
+        /*column_indices=*/{0}));
     ASSERT_EQ(5500, reader_wrapper->GetNextRowToRead());
     ASSERT_EQ(std::numeric_limits<uint64_t>::max(),
               reader_wrapper->GetPreviousBatchFirstRowNumber().value());