perf(prefetch): prefetch the late-materialization payload ranges mid-read

include/paimon/reader/prefetch_file_batch_reader.h

@@ -19,6 +19,7 @@
 
 #pragma once
 
+#include <functional>
 #include <memory>
 #include <utility>
 #include <vector>
@@ -135,6 +136,20 @@ class PAIMON_EXPORT PrefetchFileBatchReader : public FileBatchReader {
     virtual Result<std::vector<std::pair<uint64_t, uint64_t>>> PreBufferRange() {
         return std::vector<std::pair<uint64_t, uint64_t>>{};
     }
+
+    /// Callback a reader reports byte ranges through when they only become known after reading has
+    /// started, so the prefetch layer can register them with its shared read-ahead cache.
+    ///
+    /// PreBufferRange() covers what is known up front; a reader whose ranges depend on data it
+    /// has already read - the late-materialization payload pass only knows which pages hold the
+    /// matched rows once the probe pass has run - reports them through this callback instead.
+    using PreBufferRangeCallback =
+        std::function<void(std::vector<std::pair<uint64_t, uint64_t>>&&)>;
+
+    /// Installs the callback above, or clears it when `callback` is empty. By default a reader has
+    /// no late byte ranges to report and ignores the callback.
+    /// @param callback The callback to report late byte ranges through.
+    virtual void SetPreBufferRangeCallback(PreBufferRangeCallback callback) {}
 };
 
 }  // namespace paimon

include/paimon/utils/prefetch_cache_config.h

@@ -34,8 +34,9 @@ namespace paimon {
 /// ReadAheadCache to balance memory usage, I/O efficiency, and latency hiding.
 class PAIMON_EXPORT CacheConfig {
  public:
-    /// Returns the maximum allowed size (in bytes) for a single cached range.
-    /// Defaults to 32 MiB.
+    /// Returns the maximum allowed size (in bytes) for a single cached range, both for the ranges
+    /// registered up front and, as the cap, for the ones registered mid-read.
+    /// Defaults to 8 MiB.
     uint64_t GetRangeSizeLimit() const {
         return range_size_limit_;
     }
@@ -46,7 +47,7 @@ class PAIMON_EXPORT CacheConfig {
     }
 
     /// Returns the maximum gap size (in bytes) considered mergeable between
-    /// adjacent ranges. Defaults to 8 KiB.
+    /// adjacent ranges. Defaults to 512 KiB.
     uint64_t GetHoleSizeLimit() const {
         return hole_size_limit_;
     }
@@ -93,15 +94,24 @@ class PAIMON_EXPORT CacheConfig {
  private:
     // The defaults are aligned with the reader's request granularity and with
     // realistic data file sizes:
-    // - range_size_limit matches the parquet reader's 32 MiB request blocks
-    //   (Arrow ReadRangeCache's own range limit); a smaller limit cuts entries
-    //   below the request size, so a request can never be served from one piece.
+    // - range_size_limit bounds a single coalesced cached range, both for the
+    //   ranges registered up front at Init and for the ones registered mid-read.
+    //   One range is one prefetch IO, so a smaller limit spreads a large pass
+    //   over several concurrent requests instead of one long serial one; a read
+    //   spanning several adjacent ranges is still served as one hit, as they are
+    //   contiguous. 8 MiB keeps a just-in-time mid-read pass concurrent while
+    //   staying large enough to amortize each request's round trip.
     // - pre_buffer_limit must exceed the LARGEST single read a reader issues
     //   (coalesced column-chunk reads of ~128 MiB were observed): fetches are
     //   only dispatched up to this window, so a request reaching past it can
     //   never be served and falls back to a second fetch of the same bytes.
-    uint64_t range_size_limit_ = 32 * 1024 * 1024;
-    uint64_t hole_size_limit_ = 8 * 1024;
+    // - hole_size_limit trades bytes against requests: coalescing across a gap
+    //   reads the gap too, but saves a request, and on remote storage a request
+    //   costs a round trip whatever its size. The limit is therefore well above
+    //   the page-sized gaps a filtered read leaves between the pages it keeps,
+    //   which would otherwise each cost a request of their own.
+    uint64_t range_size_limit_ = 8 * 1024 * 1024;
+    uint64_t hole_size_limit_ = 512 * 1024;
     uint64_t pre_buffer_limit_ = 256 * 1024 * 1024;
     // Blocks are aligned to the END of the file, so a block never reaches past
     // EOF. 64 KiB is the granularity the reads no prefetched range covers are

src/paimon/CMakeLists.txt

@@ -902,7 +902,6 @@ if(PAIMON_BUILD_TESTS)
                     core/operation/commit/row_tracking_commit_utils_test.cpp
                     core/operation/commit/realtime_commit_properties_test.cpp
                     core/operation/commit/sequence_snapshot_properties_test.cpp
-                    core/operation/commit/retry_waiter_test.cpp
                     core/operation/key_value_file_store_write_test.cpp
                     core/operation/internal_read_context_test.cpp
                     core/operation/abstract_split_read_test.cpp

src/paimon/common/io/cache_input_stream_test.cpp

@@ -71,7 +71,7 @@ class CacheInputStreamTest : public ::testing::Test {
         // tests exercise the fallback of CacheInputStream on a cache miss.
         auto cache =
             std::make_shared<ReadAheadCache>(std::move(stream), config, /*file_size=*/0, pool_);
-        EXPECT_OK(cache->Init(std::move(ranges)));
+        EXPECT_OK(cache->AddRanges(std::move(ranges)).status());
         return cache;
     }
 
@@ -218,7 +218,7 @@ TEST_F(CacheInputStreamTest, TestReadAsyncCacheReadError) {
         config.SetPreBufferLimit(1024 * 1024);
         auto cache = std::make_shared<ReadAheadCache>(std::move(cache_stream), config,
                                                       /*file_size=*/0, pool_);
-        ASSERT_OK(cache->Init(std::vector<ByteRange>{{0, 10}}));
+        ASSERT_OK(cache->AddRanges(std::vector<ByteRange>{{0, 10}}).status());
 
         // Now activate IOHook so that the prefetch IO (triggered by cache_->Read -> PreBuffer)
         // will fail at the i-th IO operation

src/paimon/common/metrics/atomic_counter_pair.h

@@ -35,7 +35,7 @@ struct AtomicCounterPair {
     std::atomic<uint64_t> bytes{0};
 
     /// Record one request covering `size` bytes.
-    void Add(uint64_t size) {
+    void Record(uint64_t size) {
         count.fetch_add(1, std::memory_order_relaxed);
         bytes.fetch_add(size, std::memory_order_relaxed);
     }

src/paimon/common/reader/late_materializing_file_batch_reader.cpp

@@ -73,6 +73,7 @@ Result<FileBatchReader::ReadBatch> LateMaterializingFileBatchReader::NextBatch()
             PAIMON_RETURN_NOT_OK(
                 SetInnerReadSchema(payload_schema_, /*predicate=*/nullptr, matched_bitmap_));
             state_ = kRunning;
+            PAIMON_RETURN_NOT_OK(ReportPayloadPreBufferRanges());
         }
     }
 
@@ -262,6 +263,25 @@ Status LateMaterializingFileBatchReader::SetInnerReadSchema(
     return Status::OK();
 }
 
+Status LateMaterializingFileBatchReader::ReportPayloadPreBufferRanges() {
+    // The payload byte ranges only exist now: setting the payload schema refined the inner
+    // reader's target row groups down to the pages holding the matched rows. Reporting them
+    // here lets the shared read-ahead cache fetch them while this pass is still assembling
+    // its first batch, instead of every payload read waiting for its own IO.
+    if (!pre_buffer_range_callback_ || prefetch_inner_ == nullptr) {
+        return Status::OK();
+    }
+    // The error is propagated rather than swallowed: the ranges come from the file metadata, so a
+    // failure here means the reads about to follow would fail too.
+    using ByteRanges = std::vector<std::pair<uint64_t, uint64_t>>;
+    PAIMON_ASSIGN_OR_RAISE(ByteRanges ranges, prefetch_inner_->PreBufferRange());
+    if (ranges.empty()) {
+        return Status::OK();
+    }
+    pre_buffer_range_callback_(std::move(ranges));
+    return Status::OK();
+}
+
 Status LateMaterializingFileBatchReader::SetReadSchema(
     ::ArrowSchema* read_schema, const std::shared_ptr<Predicate>& predicate,
     const std::optional<RoaringBitmap32>& selection_bitmap) {

src/paimon/common/reader/late_materializing_file_batch_reader.h

@@ -93,14 +93,22 @@ class LateMaterializingFileBatchReader : public PrefetchFileBatchReader {
     Status SetReadRanges(const std::vector<std::pair<uint64_t, uint64_t>>& read_ranges) override;
 
     Result<std::vector<std::pair<uint64_t, uint64_t>>> PreBufferRange() override {
-        // TODO(zhouhongfeng.zhf): PrebufferRange (called by PrefetchFileBatchReader) only read the
-        // probe data, consider read the payload data as well.
+        // Deliberately only the probe ranges: this is asked once before any read, when the payload
+        // pass still has no idea which pages hold the matched rows. The payload ranges are
+        // reported through pre_buffer_range_callback_ once the probe pass has run, see
+        // ReportPayloadPreBufferRanges().
         if (prefetch_inner_ == nullptr) {
             return std::vector<std::pair<uint64_t, uint64_t>>{};
         }
         return prefetch_inner_->PreBufferRange();
     }
 
+    void SetPreBufferRangeCallback(PreBufferRangeCallback callback) override {
+        // Not forwarded to prefetch_inner_: a format reader has no late ranges of its own, and
+        // forwarding would report the same ranges twice.
+        pre_buffer_range_callback_ = std::move(callback);
+    }
+
  private:
     LateMaterializingFileBatchReader(std::unique_ptr<FileBatchReader> inner,
                                      PrefetchFileBatchReader* prefetch_inner,
@@ -141,6 +149,11 @@ class LateMaterializingFileBatchReader : public PrefetchFileBatchReader {
                               const std::shared_ptr<Predicate>& predicate,
                               const std::optional<RoaringBitmap32>& selection);
 
+    /// Report the byte ranges the payload pass is about to read to pre_buffer_range_callback_, so
+    /// that the prefetch layer can fetch them ahead of the pass. A no-op without a callback or
+    /// without a prefetch-capable inner reader.
+    Status ReportPayloadPreBufferRanges();
+
     /// Returns the inner reader's prefetch interface, or an error when the format reader does not
     /// implement it (avro and blob do not).
     Result<PrefetchFileBatchReader*> GetPrefetchReaderOrRaise(
@@ -163,6 +176,10 @@ class LateMaterializingFileBatchReader : public PrefetchFileBatchReader {
     /// inner_ is never reassigned, so the cast is resolved once in Create().
     PrefetchFileBatchReader* prefetch_inner_ = nullptr;
     std::shared_ptr<arrow::MemoryPool> arrow_pool_;
+    /// Where the payload byte ranges are reported. Part of how this reader is wired into the
+    /// prefetch layer rather than of a read generation, so Reset() leaves it alone.
+    /// TODO(yonghao.fyh): refactor pre buffer cache in the future
+    PreBufferRangeCallback pre_buffer_range_callback_;
     LatMatState state_ = kInit;
     std::shared_ptr<arrow::Schema> full_schema_;
     // projection holding only the predicate fields; nullptr when probing is not applicable

src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp

@@ -765,4 +765,141 @@ TEST_F(LateMaterializingFileBatchReaderTest, FailsOnPredicateTypeMismatch) {
         "mismatches");
 }
 
+// The payload byte ranges are only known once the probe pass has run, so they are reported through
+// the callback rather than through PreBufferRange(), and early enough to still be prefetched:
+// before the first payload batch leaves the reader.
+TEST_F(LateMaterializingFileBatchReaderTest, PayloadPreBufferRangesReportedToCallback) {
+    auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7, 8, 9});
+    auto mock = std::make_unique<MockFileBatchReader>(data, full_type_, /*batch_size=*/3);
+    // Aliased because a type carrying a comma cannot be passed to the ASSERT_OK_AND_ASSIGN macro.
+    using ByteRanges = std::vector<std::pair<uint64_t, uint64_t>>;
+    const ByteRanges probe_ranges = {{0, 1024}};
+    const ByteRanges payload_ranges = {{100000, 2048}, {200000, 4096}};
+    mock->SetPreBufferRangesPerSchema({probe_ranges, payload_ranges});
+    ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(
+                                          std::move(mock), GetArrowPool(pool_)));
+    int32_t batches = 0;
+    std::vector<ByteRanges> reported;
+    std::vector<int32_t> batches_when_reported;
+    reader->SetPreBufferRangeCallback([&](ByteRanges&& ranges) {
+        reported.push_back(std::move(ranges));
+        batches_when_reported.push_back(batches);
+    });
+    auto predicate =
+        PredicateBuilder::GreaterOrEqual(/*field_index=*/0, "k", FieldType::BIGINT, Literal(4l));
+    ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate, std::nullopt));
+    // What is known up front is the probe projection's ranges, not the payload's.
+    ASSERT_OK_AND_ASSIGN(ByteRanges up_front, reader->PreBufferRange());
+    EXPECT_EQ(up_front, probe_ranges);
+
+    while (true) {
+        ASSERT_OK_AND_ASSIGN(FileBatchReader::ReadBatch batch, reader->NextBatch());
+        if (BatchReader::IsEofBatch(batch)) {
+            break;
+        }
+        // Import the batch to run its release callback, as a real consumer would.
+        auto& [c_array, c_schema] = batch;
+        ASSERT_TRUE(arrow::ImportArray(c_array.get(), c_schema.get()).ok());
+        batches++;
+    }
+    EXPECT_GT(batches, 0);  // the payload pass did run
+    ASSERT_EQ(reported.size(), 1u);
+    EXPECT_EQ(reported[0], payload_ranges);
+    // Reported while the payload pass was still assembling its first batch, which is the whole
+    // point: reporting it later would leave the first payload reads waiting for their own IO.
+    EXPECT_EQ(batches_when_reported[0], 0);
+}
+
+// Without a predicate there is no payload pass, so there is no late range to report.
+TEST_F(LateMaterializingFileBatchReaderTest, NoCallbackReportWithoutLateMaterialization) {
+    auto data = BuildData({0, 1, 2, 3, 4});
+    auto mock = std::make_unique<MockFileBatchReader>(data, full_type_, /*batch_size=*/2);
+    mock->SetPreBufferRangesPerSchema({{{0, 1024}}, {{100000, 2048}}});
+    ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(
+                                          std::move(mock), GetArrowPool(pool_)));
+    int32_t calls = 0;
+    reader->SetPreBufferRangeCallback(
+        [&calls](std::vector<std::pair<uint64_t, uint64_t>>&&) { calls++; });
+    ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), /*predicate=*/nullptr,
+                            std::nullopt));
+
+    ASSERT_OK_AND_ASSIGN(std::vector<Row> rows, Collect(reader.get()));
+    ASSERT_EQ(rows.size(), 5u);
+    EXPECT_EQ(calls, 0);
+}
+
+// The probe pass matched nothing, so the payload pass never runs and reports no range.
+TEST_F(LateMaterializingFileBatchReaderTest, EmptyMatchDoesNotReportPayloadRanges) {
+    auto data = BuildData({0, 1, 2, 3, 4});
+    auto mock = std::make_unique<MockFileBatchReader>(data, full_type_, /*batch_size=*/2);
+    mock->SetPreBufferRangesPerSchema({{{0, 1024}}, {{100000, 2048}}});
+    ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(
+                                          std::move(mock), GetArrowPool(pool_)));
+    int32_t calls = 0;
+    reader->SetPreBufferRangeCallback(
+        [&calls](std::vector<std::pair<uint64_t, uint64_t>>&&) { calls++; });
+    auto predicate =
+        PredicateBuilder::GreaterThan(/*field_index=*/0, "k", FieldType::BIGINT, Literal(100l));
+    ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate, std::nullopt));
+
+    ASSERT_OK_AND_ASSIGN(std::vector<Row> rows, Collect(reader.get()));
+    EXPECT_TRUE(rows.empty());
+    EXPECT_EQ(calls, 0);
+}
+
+// The payload ranges come from the file metadata, so failing to compute them means the payload
+// reads would fail too: the error surfaces from the read instead of being swallowed.
+TEST_F(LateMaterializingFileBatchReaderTest, PayloadPreBufferRangeErrorFailsRead) {
+    auto data = BuildData({0, 1, 2, 3, 4});
+    auto mock = std::make_unique<MockFileBatchReader>(data, full_type_, /*batch_size=*/2);
+    mock->SetPreBufferRangeStatus(Status::IOError("pre-buffer range unavailable"));
+    ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(
+                                          std::move(mock), GetArrowPool(pool_)));
+    int32_t calls = 0;
+    reader->SetPreBufferRangeCallback(
+        [&calls](std::vector<std::pair<uint64_t, uint64_t>>&&) { calls++; });
+    auto predicate =
+        PredicateBuilder::GreaterOrEqual(/*field_index=*/0, "k", FieldType::BIGINT, Literal(2l));
+    ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate, std::nullopt));
+
+    ASSERT_NOK_WITH_MSG(reader->NextBatch(), "pre-buffer range unavailable");
+    EXPECT_EQ(calls, 0);
+}
+
+// End to end below the prefetch layer: the ranges the inner LM reader reports mid-read must reach
+// the shared read-ahead cache and be fetched, on top of the probe ranges registered up front.
+TEST_F(LateMaterializingFileBatchReaderTest, PrefetchInnerRegistersPayloadPreBufferRanges) {
+    auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7, 8, 9});
+    auto format_builder =
+        std::make_unique<MockFormatReaderBuilder>(data, full_type_, /*batch_size=*/3);
+    // Far enough apart that the cache's hole-size limit keeps them as two separate ranges.
+    format_builder->SetPreBufferRangesPerSchema({{{0, 1024}}, {{100000, 2048}}});
+    LateMaterializingReaderBuilder builder(std::move(format_builder), GetArrowPool(pool_));
+    auto mock_fs = std::make_shared<MockFileSystem>();
+    ASSERT_OK_AND_ASSIGN(std::shared_ptr<Executor> executor, CreateDefaultExecutor(2));
+    ASSERT_OK_AND_ASSIGN(
+        std::unique_ptr<PrefetchFileBatchReaderImpl> impl,
+        PrefetchFileBatchReaderImpl::Create(
+            /*data_file_path=*/"", /*data_file_size=*/0, &builder, mock_fs,
+            /*prefetch_max_parallel_num=*/1, /*batch_size=*/3, /*prefetch_batch_count=*/2,
+            /*enable_adaptive_prefetch_strategy=*/false, executor,
+            /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/true, CacheConfig(),
+            /*enable_io_metrics=*/false, WarmupLevel::DECODED, pool_, GetArrowPool(pool_)));
+    auto predicate =
+        PredicateBuilder::GreaterOrEqual(/*field_index=*/0, "k", FieldType::BIGINT, Literal(4l));
+    ::ArrowSchema c_schema;
+    ASSERT_TRUE(arrow::ExportSchema(*arrow::schema(full_fields_), &c_schema).ok());
+    ASSERT_OK(impl->SetReadSchema(&c_schema, predicate, std::nullopt));
+
+    ASSERT_OK_AND_ASSIGN(std::shared_ptr<arrow::StructArray> result, CollectStruct(impl.get()));
+    ASSERT_TRUE(result);
+    ASSERT_EQ(result->length(), 6);  // k = 4..9
+    std::shared_ptr<Metrics> metrics = impl->GetReaderMetrics();
+    ASSERT_TRUE(metrics);
+    ASSERT_OK_AND_ASSIGN(uint64_t io_bytes, metrics->GetCounter(ReadAheadCacheMetrics::IO_BYTES));
+    // Both the probe ranges registered before the read and the payload ranges registered during it.
+    EXPECT_EQ(io_bytes, 1024u + 2048u);
+    impl->Close();
+}
+
 }  // namespace paimon::test