diff --git a/include/paimon/reader/prefetch_file_batch_reader.h b/include/paimon/reader/prefetch_file_batch_reader.h index 842279956..1969194bf 100644 --- a/include/paimon/reader/prefetch_file_batch_reader.h +++ b/include/paimon/reader/prefetch_file_batch_reader.h @@ -19,6 +19,7 @@ #pragma once +#include #include #include #include @@ -135,6 +136,20 @@ class PAIMON_EXPORT PrefetchFileBatchReader : public FileBatchReader { virtual Result>> PreBufferRange() { return std::vector>{}; } + + /// 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>&&)>; + + /// 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 diff --git a/include/paimon/utils/prefetch_cache_config.h b/include/paimon/utils/prefetch_cache_config.h index 9f20fbf2b..d767fd929 100644 --- a/include/paimon/utils/prefetch_cache_config.h +++ b/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 diff --git a/src/paimon/CMakeLists.txt b/src/paimon/CMakeLists.txt index e32b3a9e4..f463d674a 100644 --- a/src/paimon/CMakeLists.txt +++ b/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 diff --git a/src/paimon/common/io/cache_input_stream_test.cpp b/src/paimon/common/io/cache_input_stream_test.cpp index 072af67bc..4c05c57c3 100644 --- a/src/paimon/common/io/cache_input_stream_test.cpp +++ b/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(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(std::move(cache_stream), config, /*file_size=*/0, pool_); - ASSERT_OK(cache->Init(std::vector{{0, 10}})); + ASSERT_OK(cache->AddRanges(std::vector{{0, 10}}).status()); // Now activate IOHook so that the prefetch IO (triggered by cache_->Read -> PreBuffer) // will fail at the i-th IO operation diff --git a/src/paimon/common/metrics/atomic_counter_pair.h b/src/paimon/common/metrics/atomic_counter_pair.h index 5a8297ac3..27fd17298 100644 --- a/src/paimon/common/metrics/atomic_counter_pair.h +++ b/src/paimon/common/metrics/atomic_counter_pair.h @@ -35,7 +35,7 @@ struct AtomicCounterPair { std::atomic 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); } diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp index 5554b92b2..10bdaea18 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp @@ -73,6 +73,7 @@ Result 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>; + 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, const std::optional& selection_bitmap) { diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.h b/src/paimon/common/reader/late_materializing_file_batch_reader.h index 7a225bca7..962acf6e6 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader.h +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.h @@ -93,14 +93,22 @@ class LateMaterializingFileBatchReader : public PrefetchFileBatchReader { Status SetReadRanges(const std::vector>& read_ranges) override; Result>> 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>{}; } 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 inner, PrefetchFileBatchReader* prefetch_inner, @@ -141,6 +149,11 @@ class LateMaterializingFileBatchReader : public PrefetchFileBatchReader { const std::shared_ptr& predicate, const std::optional& 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 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_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 full_schema_; // projection holding only the predicate fields; nullptr when probing is not applicable diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp b/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp index 0849836ca..3501e6190 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp +++ b/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(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>; + 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 reported; + std::vector 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(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>&&) { calls++; }); + ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), /*predicate=*/nullptr, + std::nullopt)); + + ASSERT_OK_AND_ASSIGN(std::vector 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(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>&&) { 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 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(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>&&) { 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(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(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr executor, CreateDefaultExecutor(2)); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr 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 result, CollectStruct(impl.get())); + ASSERT_TRUE(result); + ASSERT_EQ(result->length(), 6); // k = 4..9 + std::shared_ptr 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 diff --git a/src/paimon/common/reader/prefetch_file_batch_reader_impl.cpp b/src/paimon/common/reader/prefetch_file_batch_reader_impl.cpp index 2b2c7ec8b..4e5e540e4 100644 --- a/src/paimon/common/reader/prefetch_file_batch_reader_impl.cpp +++ b/src/paimon/common/reader/prefetch_file_batch_reader_impl.cpp @@ -322,6 +322,16 @@ Result> PrefetchFileBatchReaderImpl auto reader = std::unique_ptr(new PrefetchFileBatchReaderImpl( readers, batch_size, prefetch_queue_capacity, enable_adaptive_prefetch_strategy, executor, cache, io_metrics, warmup_level, arrow_pool)); + for (const auto& sub_reader : reader->readers_) { + // The sub-readers whose byte ranges only become known once they have read some data (the + // late-materialization payload pass) report them here instead of through PreBufferRange(), + // which is asked once per read-range generation before any read. Capturing the impl raw is + // safe: it owns readers_, and CleanUp() joins the background thread before it goes away. + sub_reader->SetPreBufferRangeCallback( + [impl = reader.get()](std::vector>&& ranges) { + impl->RegisterLatePreBufferRanges(std::move(ranges)); + }); + } if (initialize_read_ranges) { // normally initialize read ranges should be false, as set read schema will refresh read // ranges, and set read schema will always be called before read. @@ -356,6 +366,7 @@ PrefetchFileBatchReaderImpl::PrefetchFileBatchReaderImpl( PrefetchFileBatchReaderImpl::~PrefetchFileBatchReaderImpl() { (void)CleanUp(); + ClearPreBufferRangeCallbacks(); } Status PrefetchFileBatchReaderImpl::SetReadSchema( @@ -510,7 +521,9 @@ Status PrefetchFileBatchReaderImpl::CleanUp() { current_batch_global_row_ids_.clear(); read_ranges_freshed_ = false; // SetReadSchema()/RefreshReadRanges() call cache_->Reset() right after CleanUp(), clearing the - // cache's initialized state, so the next read-range generation must be allowed to Init again. + // cache's registered ranges, so the next read-range generation must be allowed to Init again. + // Clearing the flag also ends this generation's mid-read reports: a sub-reader reporting from + // here on registers nothing until the next generation warms the cache again. cache_warmed_.store(false); clean_prefetch_queue(); prefetch_metrics_->queue_depth.store(0, kMetricsMemoryOrder); @@ -766,10 +779,10 @@ void PrefetchFileBatchReaderImpl::WarmCacheOnce() { if (!cache_) { return; } - // Init() is not idempotent (a second call returns Invalid), so at most one warmup may run per - // read-range generation. A RAW Warmup() and the Workloop() call are ordered by the reader's own - // thread, which warms first and only then starts the background thread, so the loser of this - // exchange never continues on a cache another thread is still initializing. + // At most one warmup may run per read-range generation. A RAW Warmup() and the Workloop() call + // are ordered by the reader's own thread, which warms first and only then starts the background + // thread, so the loser of this exchange never continues on a cache another thread is still + // initializing. if (cache_warmed_.exchange(true)) { return; } @@ -782,16 +795,53 @@ void PrefetchFileBatchReaderImpl::WarmCacheOnce() { for (const auto& read_range : read_ranges.value()) { ranges.emplace_back(read_range.first, read_range.second); } - Status s = cache_->Init(std::move(ranges)); - if (!s.ok()) { - SetReadStatus(s); + Result> registered = cache_->AddRanges(std::move(ranges)); + if (!registered.ok()) { + SetReadStatus(registered.status()); return; } - // Init() only registers the ranges, so without this the first cache fetch races the readers' - // first reads instead of running ahead of them. + // AddRanges() only registers the ranges, so without this the first cache fetch races the + // readers' first reads instead of running ahead of them. cache_->Warmup(); } +void PrefetchFileBatchReaderImpl::RegisterLatePreBufferRanges( + std::vector>&& read_ranges) { + if (!cache_ || read_ranges.empty()) { + return; + } + if (!cache_warmed_.load()) { + // The read-range generation these ranges belong to has ended, so nothing is going to read + // them. WarmCacheOnce() warms the next generation. + return; + } + std::vector ranges; + ranges.reserve(read_ranges.size()); + for (const auto& read_range : read_ranges) { + ranges.emplace_back(read_range.first, read_range.second); + } + Result> first_added = cache_->AddRanges(std::move(ranges)); + if (!first_added.ok()) { + SetReadStatus(first_added.status()); + return; + } + if (!first_added.value().has_value()) { + // Every range was already registered, so there is nothing new to fetch. + return; + } + // Fetch from the first new range rather than from the start: the ranges registered before it + // belong to the pass that has already run. + cache_->Warmup(first_added.value().value()); +} + +void PrefetchFileBatchReaderImpl::ClearPreBufferRangeCallbacks() { + // Belt and braces: the callbacks capture this, and by now CleanUp() has joined the background + // thread, so no sub-reader can report a range any more. + for (const auto& sub_reader : readers_) { + sub_reader->SetPreBufferRangeCallback(nullptr); + } +} + void PrefetchFileBatchReaderImpl::Warmup() { if (warmup_level_ == WarmupLevel::NONE) { return; @@ -1032,6 +1082,7 @@ void PrefetchFileBatchReaderImpl::Close() { for (const auto& reader : readers_) { reader->Close(); } + ClearPreBufferRangeCallbacks(); } } // namespace paimon diff --git a/src/paimon/common/reader/prefetch_file_batch_reader_impl.h b/src/paimon/common/reader/prefetch_file_batch_reader_impl.h index b0b0c6f71..164cf6c6b 100644 --- a/src/paimon/common/reader/prefetch_file_batch_reader_impl.h +++ b/src/paimon/common/reader/prefetch_file_batch_reader_impl.h @@ -141,6 +141,19 @@ class PrefetchFileBatchReaderImpl : public PrefetchFileBatchReader { /// SetReadStatus() rather than returned: a warmup hint for a file that may never be read must /// not fail an in-flight read. void WarmCacheOnce(); + + /// Registers byte ranges a sub-reader discovered mid-read (the late-materialization payload + /// pass) with the shared cache and starts fetching them from there. A no-op without a cache, or + /// once the read-range generation the ranges belong to has ended (CleanUp() cleared + /// cache_warmed_): a report that outlives its generation - a cache reset by + /// SetReadSchema()/RefreshReadRanges(), or released by Close() - registers nothing instead of + /// prefetching bytes nobody reads. Errors are recorded via SetReadStatus() for the same reason + /// as in WarmCacheOnce(). + void RegisterLatePreBufferRanges(std::vector>&& read_ranges); + + /// Detaches the callbacks installed on the sub-readers, so that a sub-reader outliving this + /// reader cannot call back into it. Called once the background thread has been joined. + void ClearPreBufferRangeCallbacks(); void SetReadStatus(const Status& status); Status GetReadStatus() const; Result IsEofRange(const std::pair& read_range) const; @@ -194,7 +207,9 @@ class PrefetchFileBatchReaderImpl : public PrefetchFileBatchReader { bool need_prefetch_ = false; bool read_ranges_freshed_ = false; // Guards the one-shot read-ahead cache Init/Warmup so it runs at most once per read-range - // generation. Reset by CleanUp() alongside read_ranges_freshed_. + // generation, and gates the mid-read range reports of that generation: CleanUp() clears it, so + // a sub-reader reporting after the generation ended registers nothing. Reset by CleanUp() + // alongside read_ranges_freshed_. std::atomic cache_warmed_{false}; const uint32_t prefetch_queue_capacity_; const bool enable_adaptive_prefetch_strategy_; diff --git a/src/paimon/common/reader/prefetch_file_batch_reader_impl_test.cpp b/src/paimon/common/reader/prefetch_file_batch_reader_impl_test.cpp index aa0bd2e4e..7d0879d2a 100644 --- a/src/paimon/common/reader/prefetch_file_batch_reader_impl_test.cpp +++ b/src/paimon/common/reader/prefetch_file_batch_reader_impl_test.cpp @@ -1316,8 +1316,8 @@ TEST_P(PrefetchFileBatchReaderImplTest, TestWarmupLevelRaw) { WarmupLevel::RAW); reader->Warmup(); - // A second Warmup() must be a no-op: the one-shot guard prevents a duplicate cache Init(), - // which would otherwise fail with "Cache has already been initialized". + // A second Warmup() must be a no-op: the one-shot guard prevents opening the cache round + // twice, which would otherwise fail with "Cache has already been initialized". reader->Warmup(); std::shared_ptr metrics = reader->GetReaderMetrics(); diff --git a/src/paimon/common/utils/file_block_cache.cpp b/src/paimon/common/utils/file_block_cache.cpp index 829ffb23d..453a76eba 100644 --- a/src/paimon/common/utils/file_block_cache.cpp +++ b/src/paimon/common/utils/file_block_cache.cpp @@ -20,6 +20,8 @@ #include "paimon/common/utils/file_block_cache.h" #include +#include +#include #include "paimon/common/memory/bytes_utils.h" @@ -37,7 +39,7 @@ FileBlockCache::FileBlockCache(const std::shared_ptr& stream, uint6 FileBlockCache::~FileBlockCache() { // The fetches write into the block buffers, so they must not outlive the // stream they read from. - std::lock_guard lock(mutex_); + std::lock_guard lock(mutex_); for (auto& block : blocks_) { block.second->future.wait(); } @@ -50,11 +52,22 @@ bool FileBlockCache::Read(const ByteRange& range, char* dest) { const uint64_t index = IndexOf(range.offset); std::shared_ptr block; bool dispatch = false; + // Hit fast path under the read lock: concurrent reads of already-cached blocks + // proceed in parallel instead of serializing on an exclusive lock. { - // Publishing the promise-backed block under the lock before its fetch is - // dispatched is what makes concurrent readers of the same block wait for - // that one fetch instead of issuing their own. - std::lock_guard lock(mutex_); + std::shared_lock read_lock(mutex_); + auto it = blocks_.find(index); + if (it != blocks_.end()) { + block = it->second; + } + } + if (!block) { + // Miss: take the write lock and re-check, since another reader may have + // inserted the block while the read lock was dropped. Publishing the + // promise-backed block under the lock before its fetch is dispatched is what + // makes concurrent readers of the same block wait for that one fetch instead + // of issuing their own. + std::unique_lock write_lock(mutex_); auto it = blocks_.find(index); if (it != blocks_.end()) { block = it->second; @@ -95,32 +108,38 @@ bool FileBlockCache::Read(const ByteRange& range, char* dest) { return false; } std::memcpy(dest, block->buffer->data() + (range.offset - block->range.offset), range.length); - hits_.Add(range.length); + hit_metrics_.Record(range.length); return true; } void FileBlockCache::Release() { - std::lock_guard lock(mutex_); - // Blocks are never evicted, so waiting on blocks_ covers every dispatched - // fetch before the buffers they write into go away. - for (auto& block : blocks_) { + std::unordered_map> to_release; + { + std::lock_guard lock(mutex_); + // Swap the blocks out under the lock instead of waiting on them in place: an + // in-flight fetch may take arbitrarily long, and holding the lock while + // waiting would stall concurrent readers. + to_release.swap(blocks_); + cached_bytes_ = 0; + } + // Wait OUTSIDE the lock. Blocks are never evicted, so this covers every + // dispatched fetch before the buffers they write into go away. + for (auto& block : to_release) { block.second->future.wait(); } - blocks_.clear(); - cached_bytes_ = 0; } void FileBlockCache::ResetCounters() { - hits_.Reset(); - fetches_.Reset(); + hit_metrics_.Reset(); + fetch_metrics_.Reset(); } FileBlockCache::Counters FileBlockCache::GetCounters() const { Counters counters; - counters.hits = hits_.Count(); - counters.hit_bytes = hits_.Bytes(); - counters.fetches = fetches_.Count(); - counters.fetch_bytes = fetches_.Bytes(); + counters.hits = hit_metrics_.Count(); + counters.hit_bytes = hit_metrics_.Bytes(); + counters.fetches = fetch_metrics_.Count(); + counters.fetch_bytes = fetch_metrics_.Bytes(); return counters; } @@ -153,7 +172,7 @@ ByteRange FileBlockCache::RangeOf(uint64_t index) const { } void FileBlockCache::Fetch(const std::shared_ptr& block) { - fetches_.Add(block->range.length); + fetch_metrics_.Record(block->range.length); auto promise = block->promise; auto buffer = block->buffer; // The buffer and the promise are captured, so the async read keeps its diff --git a/src/paimon/common/utils/file_block_cache.h b/src/paimon/common/utils/file_block_cache.h index 3a2d5737a..10443b23f 100644 --- a/src/paimon/common/utils/file_block_cache.h +++ b/src/paimon/common/utils/file_block_cache.h @@ -22,7 +22,7 @@ #include #include #include -#include +#include #include #include "paimon/common/metrics/atomic_counter_pair.h" @@ -142,14 +142,14 @@ class PAIMON_EXPORT FileBlockCache { // Blocks are aligned, so keying them by index keeps them disjoint by // construction and needs no ordering. A plain mutex is enough: only the // reads that no prefetched range covers touch the map, and they are few. - mutable std::mutex mutex_; + mutable std::shared_mutex mutex_; std::unordered_map> blocks_; // Bytes held by blocks_, guarded by mutex_ and bounded by capacity_. uint64_t cached_bytes_ = 0; // The requests served out of a block, and the block fetches issued to the // underlying stream. - AtomicCounterPair hits_; - AtomicCounterPair fetches_; + AtomicCounterPair hit_metrics_; + AtomicCounterPair fetch_metrics_; }; } // namespace paimon diff --git a/src/paimon/common/utils/read_ahead_cache.cpp b/src/paimon/common/utils/read_ahead_cache.cpp index 6e517aaaf..68dc18034 100644 --- a/src/paimon/common/utils/read_ahead_cache.cpp +++ b/src/paimon/common/utils/read_ahead_cache.cpp @@ -23,11 +23,12 @@ #include "paimon/common/utils/read_ahead_cache.h" #include -#include #include #include #include +#include #include +#include #include "paimon/common/memory/bytes_utils.h" #include "paimon/common/metrics/atomic_counter_pair.h" @@ -39,23 +40,8 @@ namespace paimon { -struct RangeCacheEntry { - ByteRange range; - std::shared_ptr buffer; - std::shared_future future; // use shared_future in case of multiple get calls - - RangeCacheEntry() = default; - RangeCacheEntry(const ByteRange& range, std::shared_ptr buffer, - std::future future) - : range(range), buffer(std::move(buffer)), future(std::move(future).share()) {} - - friend bool operator<(const RangeCacheEntry& left, const RangeCacheEntry& right) { - return left.range.offset < right.range.offset; - } -}; - // Everything needed to dispatch the prefetch IO of an entry AFTER the entry -// has been published into entries_: the promise resolves the entry's future +// has been published into ranges_: the promise resolves the entry's future // and the buffer capture keeps the destination alive for the async IO. struct PendingFetch { ByteRange range; @@ -82,146 +68,154 @@ void CopyRangeFromEntries(const std::vector& covering, const By } // namespace -class ReadAheadCache::Impl { - public: - Impl(const std::shared_ptr& stream, const CacheConfig& config, uint64_t file_size, - const std::shared_ptr& memory_pool); - ~Impl(); - - Status Init(std::vector&& ranges); - Result Read(const ByteRange& range, char* dest); - void Reset(); - void ReleaseBuffers(); - void Warmup(); - void CollectMetrics(std::shared_ptr* metrics) const; - - private: - /// Dispatch the prefetch IOs for entries that have already been published - /// into entries_. - void DispatchFetches(const std::vector& fetches); - /// Find the entries fully covering the given range under the read lock. - /// Returns an empty vector on miss. Entries are copied (shared buffers) - /// so the caller may use them after releasing the lock. - std::vector FindCoveringEntries(const ByteRange& range); - void PreBuffer(uint64_t offset); - - /// Mark, publish and fetch the pending ranges at the given indices. - /// - /// Marking is_cached_ and publishing the promise-backed entries happen - /// atomically under the write lock, before any IO is dispatched, so a - /// reader racing the prefetch waits on the in-flight entries instead of - /// re-fetching the same bytes. - void Cache(std::vector pending_indices); - - /// Clear the prefetch state, waiting for the fetches still writing into the - /// entry buffers. Leaves the block cache untouched. - void ReleasePrefetchBuffers(); +Result> ReadAheadCache::AddRanges(std::vector&& ranges) { + // Both an up-front registration and a mid-read one are cut at the plain size limit. + PAIMON_ASSIGN_OR_RAISE( + std::vector new_ranges, + ByteRangeCombiner::CoalesceByteRanges(std::move(ranges), config_.GetHoleSizeLimit(), + config_.GetRangeSizeLimit())); + for (const auto& new_range : new_ranges) { + PAIMON_RETURN_NOT_OK(ValidateValueInRange(new_range.offset, "range offset")); + PAIMON_RETURN_NOT_OK(ValidateValueInRange(new_range.length, "range length")); + } - std::shared_ptr stream_; - CacheConfig config_; - // Ordered by offset (so as to find a matching region by binary search) - std::vector entries_; - std::shared_ptr memory_pool_; - std::shared_mutex rw_mutex_; - std::vector> is_cached_; - std::vector pending_ranges_; - bool is_initialized_ = false; - // Caches the reads that no registered range covers, or null when the block - // cache is disabled. Owns its own locking and counters. - std::unique_ptr block_cache_; - // The Read() requests issued to the cache and how they were served, - // aggregated over all the streams sharing this cache. A read is counted - // either as a hit, a block cache hit or a miss. - AtomicCounterPair reads_; - AtomicCounterPair hits_; - AtomicCounterPair misses_; - // The prefetch IO actually issued to the underlying stream. The block cache - // counts its own fetches, which CollectMetrics() adds to these. - AtomicCounterPair ios_; -}; + // Locked: a registration rewrites the very ranges_ a concurrent Read() serves from, so it + // must not race the read finding its covering entries. + std::unique_lock lock(rw_mutex_); + return MergeRangesLocked(std::move(new_ranges)); +} -void ReadAheadCache::Impl::Cache(std::vector pending_indices) { - std::vector new_entries; - std::vector fetches; - // Mark is_cached_, publish the promise-backed entries and only then - // dispatch the IOs. The mark and the publication happen atomically under - // the write lock: a reader racing the prefetch observes is_cached_=true - // only once the covering entries are already visible, so it waits on - // their futures instead of issuing a duplicate underlying read. - { - std::unique_lock lock(rw_mutex_); - for (size_t idx : pending_indices) { - if (is_cached_[idx].exchange(true)) { +std::optional ReadAheadCache::MergeRangesLocked(std::vector&& new_ranges) { + // Merge the two sorted, internally disjoint lists, keeping the result disjoint so that + // FindCoveringEntries() may keep walking it: of a new range, only the parts no registered + // range covers are registered. Already registered entries are moved over as-is, so a range + // that was published keeps its buffer and future and is never fetched again. + std::vector merged; + merged.reserve(ranges_.size() + new_ranges.size()); + std::optional first_added; + size_t old_idx = 0; + auto keep_registered = [&]() { + merged.push_back(std::move(ranges_[old_idx])); + ++old_idx; + }; + auto register_new = [&](uint64_t offset, uint64_t length) { + if (!first_added.has_value()) { + first_added = offset; + } + merged.emplace_back(ByteRange(offset, length)); + }; + for (const ByteRange& candidate : new_ranges) { + const uint64_t candidate_end = candidate.offset + candidate.length; + // Start of the part of the candidate that is not registered yet. + uint64_t cursor = candidate.offset; + while (cursor < candidate_end) { + if (old_idx == ranges_.size()) { + register_new(cursor, candidate_end - cursor); + break; + } + const ByteRange registered = ranges_[old_idx].range; + const uint64_t registered_end = registered.offset + registered.length; + if (registered_end <= cursor) { + // Entirely before what is left of the candidate, so its place in the merged list + // is settled. + keep_registered(); continue; } - const ByteRange& range = pending_ranges_[idx]; - auto promise = std::make_shared>(); - auto future = promise->get_future(); - auto buffer = AllocateBytesKeepingPoolAlive(range.length, memory_pool_); - fetches.push_back({range, buffer, promise}); - new_entries.emplace_back(range, std::move(buffer), std::move(future)); - } - if (!new_entries.empty()) { - // Entries are never evicted: the cache holds every published - // range until ReleaseBuffers()/Reset(), so an in-flight fetch - // always keeps its entry and thus its future reachable. - std::vector merged(entries_.size() + new_entries.size()); - std::merge(entries_.begin(), entries_.end(), new_entries.begin(), new_entries.end(), - merged.begin()); - entries_ = std::move(merged); + if (registered.offset >= candidate_end) { + // The candidate ends before this one starts: the rest of it is new. + register_new(cursor, candidate_end - cursor); + break; + } + // Overlap: the hole before the registered range is new, and the overlapping part is + // dropped because it is already being fetched. The registered range is left for the + // next iteration to settle, as a later candidate may still start before it. + if (registered.offset > cursor) { + register_new(cursor, registered.offset - cursor); + } + cursor = registered_end; } } - DispatchFetches(fetches); -} - -Status ReadAheadCache::Impl::Init(std::vector&& ranges) { - if (is_initialized_) { - return Status::Invalid("Cache has already been initialized"); - } - PAIMON_ASSIGN_OR_RAISE( - std::vector pending_ranges, - ByteRangeCombiner::CoalesceByteRanges(std::move(ranges), config_.GetHoleSizeLimit(), - config_.GetRangeSizeLimit())); - for (const auto& pending_range : pending_ranges) { - PAIMON_RETURN_NOT_OK(ValidateValueInRange(pending_range.offset, "range offset")); - PAIMON_RETURN_NOT_OK(ValidateValueInRange(pending_range.length, "range length")); - } - pending_ranges_ = pending_ranges; - is_cached_ = std::vector>(pending_ranges_.size()); - for (auto& is_cached : is_cached_) { - is_cached.store(false); + while (old_idx < ranges_.size()) { + keep_registered(); } - is_initialized_ = true; - return Status::OK(); + // Always move the merged list back: keep_registered() has already moved the existing entries' + // buffer/future out of ranges_, so returning early without this would drop the published state + // (and the future to wait on) even when nothing new was registered. + ranges_ = std::move(merged); + return first_added; } -void ReadAheadCache::Impl::PreBuffer(uint64_t offset) { - auto it = std::lower_bound(pending_ranges_.begin(), pending_ranges_.end(), offset, - [](const ByteRange& range, uint64_t offset) { - return range.offset + range.length <= offset; +bool ReadAheadCache::WindowHasUnpublishedLocked(uint64_t offset) const { + auto it = std::lower_bound(ranges_.begin(), ranges_.end(), offset, + [](const RangeCacheEntry& entry, uint64_t offset) { + return entry.range.offset + entry.range.length <= offset; }); - if (it == pending_ranges_.end() || it->offset > offset) { - return; + if (it == ranges_.end() || it->range.offset > offset) { + return false; } - - size_t start_idx = std::distance(pending_ranges_.begin(), it); - std::vector pending_indices; size_t total_bytes = 0; - for (size_t i = start_idx; i < pending_ranges_.size(); ++i) { - total_bytes += pending_ranges_[i].length; + for (auto entry = it; entry != ranges_.end(); ++entry) { + total_bytes += entry->range.length; if (total_bytes > config_.GetPreBufferLimit()) { break; } - pending_indices.push_back(i); + if (!entry->Published()) { + return true; + } } + return false; +} - if (!pending_indices.empty()) { - Cache(std::move(pending_indices)); +void ReadAheadCache::PreBuffer(uint64_t offset) { + // Fast path under the read lock: once the whole prefetch window is published there is nothing + // left to fetch, and returning here keeps the hot Read() path off the write lock, so concurrent + // readers' lookups are not serialized behind a publish that has nothing to publish. + { + std::shared_lock read_lock(rw_mutex_); + if (!WindowHasUnpublishedLocked(offset)) { + return; + } } + std::vector fetches; + { + std::unique_lock lock(rw_mutex_); + auto it = std::lower_bound(ranges_.begin(), ranges_.end(), offset, + [](const RangeCacheEntry& entry, uint64_t offset) { + return entry.range.offset + entry.range.length <= offset; + }); + if (it == ranges_.end() || it->range.offset > offset) { + return; + } + + size_t total_bytes = 0; + for (auto entry = it; entry != ranges_.end(); ++entry) { + total_bytes += entry->range.length; + if (total_bytes > config_.GetPreBufferLimit()) { + break; + } + if (entry->Published()) { + continue; + } + auto promise = std::make_shared>(); + auto future = promise->get_future(); + auto buffer = AllocateBytesKeepingPoolAlive(entry->range.length, memory_pool_); + fetches.push_back({entry->range, buffer, promise}); + // Published IN PLACE under the write lock, before the IO is dispatched below: a + // reader racing the prefetch then finds the entry and waits on its future instead of + // issuing a second fetch for the same bytes. Entries are never evicted, so an + // in-flight fetch always keeps its entry and thus its future reachable. + entry->buffer = std::move(buffer); + entry->future = std::move(future).share(); + } + } + // Dispatched OUTSIDE the lock, so that a fetch completing inline does not deadlock against a + // Read() waiting for it. + DispatchFetches(fetches); } -ReadAheadCache::Impl::Impl(const std::shared_ptr& stream, const CacheConfig& config, - uint64_t file_size, const std::shared_ptr& memory_pool) +ReadAheadCache::ReadAheadCache(const std::shared_ptr& stream, + const CacheConfig& config, uint64_t file_size, + const std::shared_ptr& memory_pool) : stream_(stream), config_(config), memory_pool_(memory_pool) { // An unknown file size cannot be aligned to, and a zero limit or block size // means the block cache is turned off: leave it null in those cases. @@ -231,20 +225,22 @@ ReadAheadCache::Impl::Impl(const std::shared_ptr& stream, const Cac } } -ReadAheadCache::Impl::~Impl() { +ReadAheadCache::~ReadAheadCache() { std::unique_lock lock(rw_mutex_); - for (auto& entry : entries_) { - entry.future.wait(); + for (auto& entry : ranges_) { + if (entry.Published()) { + entry.future.wait(); + } } // The block cache waits for its own fetches when it is destroyed. } -void ReadAheadCache::Impl::Reset() { +void ReadAheadCache::Reset() { ReleasePrefetchBuffers(); - reads_.Reset(); - hits_.Reset(); - misses_.Reset(); - ios_.Reset(); + read_metrics_.Reset(); + hit_metrics_.Reset(); + miss_metrics_.Reset(); + io_metrics_.Reset(); if (block_cache_ != nullptr) { // Only the counters: the blocks cache the file, not the registered // ranges, and a reader resetting the cache reads the same file again. @@ -252,41 +248,47 @@ void ReadAheadCache::Impl::Reset() { } } -void ReadAheadCache::Impl::ReleaseBuffers() { +void ReadAheadCache::ReleaseBuffers() { ReleasePrefetchBuffers(); if (block_cache_ != nullptr) { block_cache_->Release(); } } -void ReadAheadCache::Impl::ReleasePrefetchBuffers() { - std::unique_lock lock(rw_mutex_); - // Entries are never evicted, so waiting on entries_ covers every - // dispatched fetch: no fetch is still writing into an entry buffer when the - // buffers go away. The buffers keep the memory pool alive themselves, for - // the callbacks that a stream destroys later than it resolves them. - for (auto& entry : entries_) { - entry.future.wait(); +void ReadAheadCache::ReleasePrefetchBuffers() { + std::vector to_release; + { + std::unique_lock lock(rw_mutex_); + // Swap the entries out under the lock instead of waiting on them in place: an in-flight + // fetch may take arbitrarily long, and holding the write lock while waiting would stall + // every concurrent reader. Once swapped, the entries are owned locally and no other thread + // can observe or mutate them. Entries are never evicted, so waiting on the published ones + // covers every dispatched fetch: no fetch is still writing into an entry buffer when the + // buffers go away. The buffers keep the memory pool alive themselves, for the callbacks + // that a stream destroys later than it resolves them. + to_release.swap(ranges_); + } + // Wait OUTSIDE the lock; the buffers are released when to_release destructs. + for (auto& entry : to_release) { + if (entry.Published()) { + entry.future.wait(); + } } - entries_.clear(); - is_cached_.clear(); - pending_ranges_.clear(); - is_initialized_ = false; // The read/io counters are deliberately kept: a reader closed at EOF must // still be able to report them through CollectMetrics(). } -void ReadAheadCache::Impl::CollectMetrics(std::shared_ptr* metrics) const { +void ReadAheadCache::CollectMetrics(std::shared_ptr* metrics) const { if (metrics == nullptr || !*metrics) { return; } auto& m = *metrics; - m->SetCounter(ReadAheadCacheMetrics::READ_COUNT, reads_.Count()); - m->SetCounter(ReadAheadCacheMetrics::READ_BYTES, reads_.Bytes()); - m->SetCounter(ReadAheadCacheMetrics::READ_HITS, hits_.Count()); - m->SetCounter(ReadAheadCacheMetrics::READ_HIT_BYTES, hits_.Bytes()); - m->SetCounter(ReadAheadCacheMetrics::READ_MISSES, misses_.Count()); - m->SetCounter(ReadAheadCacheMetrics::READ_MISS_BYTES, misses_.Bytes()); + m->SetCounter(ReadAheadCacheMetrics::READ_COUNT, read_metrics_.Count()); + m->SetCounter(ReadAheadCacheMetrics::READ_BYTES, read_metrics_.Bytes()); + m->SetCounter(ReadAheadCacheMetrics::READ_HITS, hit_metrics_.Count()); + m->SetCounter(ReadAheadCacheMetrics::READ_HIT_BYTES, hit_metrics_.Bytes()); + m->SetCounter(ReadAheadCacheMetrics::READ_MISSES, miss_metrics_.Count()); + m->SetCounter(ReadAheadCacheMetrics::READ_MISS_BYTES, miss_metrics_.Bytes()); // The block cache keeps its own counters. Its fetches also go to the // underlying stream, so they are part of the io counters too. const FileBlockCache::Counters blocks = @@ -295,28 +297,40 @@ void ReadAheadCache::Impl::CollectMetrics(std::shared_ptr* metrics) con m->SetCounter(ReadAheadCacheMetrics::BLOCK_HIT_BYTES, blocks.hit_bytes); m->SetCounter(ReadAheadCacheMetrics::BLOCK_FETCHES, blocks.fetches); m->SetCounter(ReadAheadCacheMetrics::BLOCK_FETCH_BYTES, blocks.fetch_bytes); - m->SetCounter(ReadAheadCacheMetrics::IO_COUNT, ios_.Count() + blocks.fetches); - m->SetCounter(ReadAheadCacheMetrics::IO_BYTES, ios_.Bytes() + blocks.fetch_bytes); + m->SetCounter(ReadAheadCacheMetrics::IO_COUNT, io_metrics_.Count() + blocks.fetches); + m->SetCounter(ReadAheadCacheMetrics::IO_BYTES, io_metrics_.Bytes() + blocks.fetch_bytes); } -void ReadAheadCache::Impl::Warmup() { - // Init() only registers the pending ranges; without this the first fetch +void ReadAheadCache::Warmup() { + // AddRanges() only registers the pending ranges; without this the first fetch // starts when the first Read() arrives, racing the reader's own miss fetch. - if (!pending_ranges_.empty()) { - PreBuffer(pending_ranges_.front().offset); + uint64_t from_offset = 0; + { + std::shared_lock lock(rw_mutex_); + if (ranges_.empty()) { + return; + } + from_offset = ranges_.front().range.offset; } + PreBuffer(from_offset); +} + +void ReadAheadCache::Warmup(uint64_t from_offset) { + PreBuffer(from_offset); } -std::vector ReadAheadCache::Impl::FindCoveringEntries(const ByteRange& range) { +std::vector ReadAheadCache::FindCoveringEntries(const ByteRange& range) { std::vector covering; std::shared_lock lock(rw_mutex_); // Find the entry holding the start of the range: the first entry whose // end is beyond range.offset (entries are disjoint and sorted by offset). - auto it = std::lower_bound(entries_.begin(), entries_.end(), range.offset, + auto it = std::lower_bound(ranges_.begin(), ranges_.end(), range.offset, [](const RangeCacheEntry& e, uint64_t offset) { return e.range.offset + e.range.length <= offset; }); - if (it == entries_.end() || it->range.offset > range.offset) { + // An unpublished covering entry means its prefetch has not been dispatched + // yet, so the bytes are not there to serve: a miss. + if (it == ranges_.end() || it->range.offset > range.offset || !it->Published()) { return covering; } if (it->range.Contains(range)) { @@ -324,15 +338,15 @@ std::vector ReadAheadCache::Impl::FindCoveringEntries(const Byt return covering; } // The request spans several adjacent entries (a column chunk larger than - // one coalesced range): collect the contiguous run and check it covers - // the whole request. Entries are published before their fetch is - // dispatched, so a reader racing the prefetch waits for the in-flight - // fetch instead of issuing a second one for the same bytes. + // one coalesced range): collect the contiguous run of published entries and + // check it covers the whole request. Entries are published before their + // fetch is dispatched, so a reader racing the prefetch waits for the + // in-flight fetch instead of issuing a second one for the same bytes. uint64_t covered_end = it->range.offset + it->range.length; covering.push_back(*it); auto next = std::next(it); - while (covered_end < range.offset + range.length && next != entries_.end() && - next->range.offset == covered_end) { + while (covered_end < range.offset + range.length && next != ranges_.end() && + next->range.offset == covered_end && next->Published()) { covered_end = next->range.offset + next->range.length; covering.push_back(*next); ++next; @@ -343,11 +357,11 @@ std::vector ReadAheadCache::Impl::FindCoveringEntries(const Byt return covering; } -Result ReadAheadCache::Impl::Read(const ByteRange& range, char* dest) { +Result ReadAheadCache::Read(const ByteRange& range, char* dest) { if (range.length == 0) { return true; } - reads_.Add(range.length); + read_metrics_.Record(range.length); PreBuffer(range.offset); std::vector covering = FindCoveringEntries(range); if (covering.empty()) { @@ -358,7 +372,7 @@ Result ReadAheadCache::Impl::Read(const ByteRange& range, char* dest) { // The block cache counts its own hits, see CollectMetrics(). return true; } - misses_.Add(range.length); + miss_metrics_.Record(range.length); return false; } // Wait OUTSIDE the lock: the futures resolve when the prefetch stream's @@ -368,11 +382,11 @@ Result ReadAheadCache::Impl::Read(const ByteRange& range, char* dest) { } // The data copy runs OUTSIDE the lock for the same reason. CopyRangeFromEntries(covering, range, dest); - hits_.Add(range.length); + hit_metrics_.Record(range.length); return true; } -void ReadAheadCache::Impl::DispatchFetches(const std::vector& fetches) { +void ReadAheadCache::DispatchFetches(const std::vector& fetches) { for (const auto& fetch : fetches) { auto promise = fetch.promise; auto buffer = fetch.buffer; @@ -381,39 +395,8 @@ void ReadAheadCache::Impl::DispatchFetches(const std::vector& fetc stream_->ReadAsync( buffer->data(), read_size, read_offset, [promise, buffer](Status status) mutable { promise->set_value(status); }); - ios_.Add(fetch.range.length); + io_metrics_.Record(fetch.range.length); } } -ReadAheadCache::ReadAheadCache(const std::shared_ptr& stream, - const CacheConfig& config, uint64_t file_size, - const std::shared_ptr& memory_pool) - : impl_(std::make_unique(stream, config, file_size, memory_pool)) {} - -ReadAheadCache::~ReadAheadCache() = default; - -Status ReadAheadCache::Init(std::vector&& ranges) { - return impl_->Init(std::move(ranges)); -} - -Result ReadAheadCache::Read(const ByteRange& range, char* dest) { - return impl_->Read(range, dest); -} - -void ReadAheadCache::Reset() { - return impl_->Reset(); -} - -void ReadAheadCache::ReleaseBuffers() { - return impl_->ReleaseBuffers(); -} - -void ReadAheadCache::Warmup() { - impl_->Warmup(); -} - -void ReadAheadCache::CollectMetrics(std::shared_ptr* metrics) const { - impl_->CollectMetrics(metrics); -} - } // namespace paimon diff --git a/src/paimon/common/utils/read_ahead_cache.h b/src/paimon/common/utils/read_ahead_cache.h index e44c277d2..4eec73836 100644 --- a/src/paimon/common/utils/read_ahead_cache.h +++ b/src/paimon/common/utils/read_ahead_cache.h @@ -23,9 +23,14 @@ #pragma once #include +#include #include +#include +#include +#include #include +#include "paimon/common/metrics/atomic_counter_pair.h" #include "paimon/fs/file_system.h" #include "paimon/memory/memory_pool.h" #include "paimon/result.h" @@ -36,6 +41,9 @@ namespace paimon { class Metrics; +class FileBlockCache; +class Bytes; +struct PendingFetch; /// Metric names for the read-ahead cache. class PAIMON_EXPORT ReadAheadCacheMetrics { @@ -98,11 +106,36 @@ struct PAIMON_EXPORT ByteRange { } }; -/// A read cache designed to hide IO latencies when reading. -/// Prefetching strategy: When a range is read, the cache will prefetch up to -/// `pre_buffer_range_count` additional adjacent ranges ahead of the requested offset. This helps -/// hide I/O latency for sequential access. Example: If you read range [0, 100), and -/// pre_buffer_range_count=2, the next two configured ranges will also be prefetched. +/// A registered range of the cache. `buffer` is null until the prefetch for the +/// range is published (dispatched); once published it holds the destination the +/// async IO writes into and `future` resolves when that IO completes. So +/// `Published()` distinguishes "registered, not fetched yet" from "fetched". +struct RangeCacheEntry { + ByteRange range; + std::shared_ptr buffer; + std::shared_future future; // use shared_future in case of multiple get calls + + RangeCacheEntry() = default; + /// Register a range without publishing it: no buffer, no fetch yet. + explicit RangeCacheEntry(const ByteRange& range) : range(range) {} + RangeCacheEntry(const ByteRange& range, std::shared_ptr buffer, + std::future future) + : range(range), buffer(std::move(buffer)), future(std::move(future).share()) {} + + /// True once the prefetch for this range has been published (buffer assigned). + bool Published() const { + return buffer != nullptr; + } + + friend bool operator<(const RangeCacheEntry& left, const RangeCacheEntry& right) { + return left.range.offset < right.range.offset; + } +}; + +/// A read cache designed to hide IO latency when reading. +/// Prefetching strategy: a Read() (or Warmup()) at an offset publishes and fetches the registered +/// ranges from that offset forward until their cumulative size exceeds `pre_buffer_limit`, so a +/// sequential read runs ahead of the caller instead of fetching one range at a time. /// /// The cache never evicts: every published range stays cached until /// ReleaseBuffers() or Reset(). It is meant to hold the prefetched ranges of @@ -113,7 +146,7 @@ struct PAIMON_EXPORT ByteRange { /// instead of being left to the caller. That block cache is owned by this one /// and shares its lifetime: it is configured from `block_size` and /// `block_cache_limit`, it survives Reset() - the blocks belong to the file -/// rather than to a registration round - and it is released by ReleaseBuffers(). +/// rather than to the registered ranges - and it is released by ReleaseBuffers(). class PAIMON_EXPORT ReadAheadCache { public: /// Construct a read cache with given options @@ -127,14 +160,23 @@ class PAIMON_EXPORT ReadAheadCache { uint64_t file_size, const std::shared_ptr& memory_pool); ~ReadAheadCache(); - /// Initialize the cache with given byte ranges to be cached. - /// @param ranges The byte ranges to be cached. - /// @return Status of the operation. - /// @note This method must be called before any Read() calls. Ranges will be coalesced based - /// on the cache configuration. - Status Init(std::vector&& ranges); + /// Register byte ranges to prefetch, merging them into the ranges already registered. + /// + /// The cache's only registration entry point; safe to call repeatedly and concurrently with + /// Read(). Ranges known up front are registered before any Read(); ranges that only become + /// known while reading (a late-materialization payload pass learns its pages after the probe + /// pass) are registered as they are discovered. The list stays disjoint and offset-ordered, so + /// only the part of a new range that no registered range covers is added - the overlap is + /// dropped, as those bytes are already being fetched - and each registration is cut into ranges + /// of at most `range_size_limit` bytes (one prefetch IO each) so a large pass is fetched + /// concurrently. Which ranges to register, and not reporting those of an ended read, is the + /// caller's concern. + /// + /// @param ranges The byte ranges to register. + /// @return The offset of the first newly registered range, or nullopt when nothing was added. + Result> AddRanges(std::vector&& ranges); - /// Read a range previously provided to Init(), copying the cached data + /// Read a range previously registered through AddRanges(), copying the cached data /// directly into the given destination buffer. /// /// Multi-segment hits are copied into `dest` segment by segment, without @@ -148,31 +190,33 @@ class PAIMON_EXPORT ReadAheadCache { /// filled; false on cache miss (`dest` is left untouched). Result Read(const ByteRange& range, char* dest); - /// Start fetching the first batch of pending ranges immediately. - /// Init() only registers the ranges; without Warmup() the first fetch starts - /// when the first Read() arrives, racing the caller's own miss fetch. + /// Start fetching from the first registered range forward immediately. AddRanges() only + /// registers the ranges; without Warmup() the first fetch starts when the first Read() arrives, + /// racing the caller's own miss fetch. void Warmup(); - /// Collect the counters of the Read() calls, of the block cache and of the - /// IOs into the given metrics as counters named after - /// `ReadAheadCacheMetrics`. Only reads issued through Read() are counted - /// as hits, block hits or misses; the fetches the cache dispatches itself - /// are counted in the io counters instead. + /// Start fetching the registered ranges from `from_offset` forward, bounded by the + /// pre-buffer limit. Warmup() is this call made from the first registered range. + /// @param from_offset The offset to start fetching from, typically one AddRanges() returned. + void Warmup(uint64_t from_offset); + + /// Collect the counters of the Read() calls, of the block cache and of the IOs into the given + /// metrics as counters named after `ReadAheadCacheMetrics`. Only reads issued through Read() + /// are counted as hits, block hits or misses; the fetches the cache dispatches are counted in + /// the io counters instead. /// @param metrics The metrics to write the counters into. A null /// pointer or a null shared pointer is a no-op. void CollectMetrics(std::shared_ptr* metrics) const; - /// Reset the cache to its initial state, clearing all cached data and configuration. - /// - /// This method waits for all ongoing asynchronous read operations to complete, - /// clears all cached entries, and resets the internal state so that Init() can be called again. - /// After calling Reset, the cache can be safely re-initialized with new ranges. + /// Drop the prefetched entries and zero all counters, so a fresh set of ranges can be + /// registered with AddRanges(). Waits for the in-flight fetches before their buffers go away. /// - /// The block cache is kept: it caches the file rather than the registered - /// ranges, and a reader reusing the cache reads the same file again. + /// The block cache is kept - it caches the file rather than the registered ranges, and a reader + /// reusing the cache reads the same file again - but its counters are zeroed too. Unlike + /// ReleaseBuffers(), the counters do not survive. void Reset(); - /// Release all cached buffers and pending ranges while keeping the hit/miss + /// Release all cached buffers and registered ranges while keeping the hit/miss /// counters intact. /// /// Unlike Reset(), the counters recorded by Read() remain readable through @@ -182,8 +226,59 @@ class PAIMON_EXPORT ReadAheadCache { void ReleaseBuffers(); private: - class Impl; - std::unique_ptr impl_; + /// Merge coalesced, validated `new_ranges` into the disjoint, offset-ordered ranges_, moving + /// the already registered entries (and thus their published buffer/future) over as-is, and + /// return the offset of the first newly registered range (nullopt when nothing was added). + /// The caller holds the write lock. + std::optional MergeRangesLocked(std::vector&& new_ranges); + + /// Dispatch the prefetch IOs for entries that have already been published into ranges_. + void DispatchFetches(const std::vector& fetches); + + /// Find the published entries fully covering the given range under the read lock. + /// Returns an empty vector on miss. Entries are copied (shared buffers) + /// so the caller may use them after releasing the lock. + std::vector FindCoveringEntries(const ByteRange& range); + + /// True when the prefetch window starting at `offset` still holds an entry whose fetch has not + /// been published. The caller holds either lock; PreBuffer() uses it under the read lock to + /// skip the write lock entirely when there is nothing left to fetch. + bool WindowHasUnpublishedLocked(uint64_t offset) const; + + /// Publish and fetch the registered ranges from the given offset forward, bounded by the + /// pre-buffer limit. + /// + /// Selecting the ranges and publishing them (filling buffer/future in place in ranges_) all + /// happen under the write lock, before any IO is dispatched: AddRanges() rewrites ranges_ + /// from another thread, and a reader racing the prefetch must observe an entry published only + /// once it is already visible in ranges_, so it waits on the in-flight entry instead of + /// re-fetching the same bytes. + void PreBuffer(uint64_t offset); + + /// Clear the prefetch state, waiting for the fetches still writing into the + /// entry buffers. Leaves the block cache untouched. + void ReleasePrefetchBuffers(); + + std::shared_ptr stream_; + CacheConfig config_; + // Every registered range, ordered by offset and disjoint (so a matching + // region is found by binary search). An entry whose buffer is still null is + // registered but not fetched yet; publishing fills buffer/future in place. + std::vector ranges_; + std::shared_ptr memory_pool_; + mutable std::shared_mutex rw_mutex_; + // Caches the reads that no registered range covers, or null when the block + // cache is disabled. Owns its own locking and counters. + std::unique_ptr block_cache_; + // The Read() requests issued to the cache and how they were served, + // aggregated over all the streams sharing this cache. A read is counted + // either as a hit, a block cache hit or a miss. + AtomicCounterPair read_metrics_; + AtomicCounterPair hit_metrics_; + AtomicCounterPair miss_metrics_; + // The prefetch IO actually issued to the underlying stream. The block cache + // counts its own fetches, which CollectMetrics() adds to these. + AtomicCounterPair io_metrics_; }; } // namespace paimon diff --git a/src/paimon/common/utils/read_ahead_cache_test.cpp b/src/paimon/common/utils/read_ahead_cache_test.cpp index 60fe87094..c9f7454a1 100644 --- a/src/paimon/common/utils/read_ahead_cache_test.cpp +++ b/src/paimon/common/utils/read_ahead_cache_test.cpp @@ -22,6 +22,7 @@ #include #include #include +#include #include #include @@ -37,7 +38,9 @@ namespace paimon::test { // The range limits the tests exercise. The defaults are sized for real data -// files, which the small test files would never reach. +// files, which the small test files would never reach. The ranges registered +// mid-read are cut at the same limit as the ones registered up front, so a test +// that does not care about the late granularity sees one range per registration. CacheConfig TestCacheConfig(uint64_t range_size_limit, uint64_t hole_size_limit, uint64_t pre_buffer_limit) { CacheConfig config; @@ -85,7 +88,7 @@ std::shared_ptr CreateTestFileAndCache(const std::string& filena std::shared_ptr in = OpenTestFile(&dir, filename, content); std::shared_ptr cache = std::make_shared(in, config, file_size, TestPool()); - EXPECT_OK(cache->Init(std::move(ranges))); + EXPECT_OK(cache->AddRanges(std::move(ranges)).status()); return cache; } @@ -108,6 +111,13 @@ void AssertReadMiss(const ByteRange& range, ReadAheadCache* cache) { ASSERT_EQ(std::string(dest.size(), 'X'), dest); } +// Report the counter of the given metric name, or zero when it was not collected. +uint64_t CounterOf(const char* name, ReadAheadCache* cache) { + std::shared_ptr metrics = std::make_shared(); + cache->CollectMetrics(&metrics); + return metrics->GetCounter(name).value_or(0); +} + TEST(TestReadAheadCache, TestBasics) { CacheConfig config = TestCacheConfig(/*range_size_limit=*/10, @@ -142,7 +152,7 @@ TEST(TestReadAheadCache, TestMultiSegmentContiguousHit) { CacheConfig config = TestCacheConfig(/*range_size_limit=*/10, /*hole_size_limit=*/2, /*pre_buffer_limit=*/1024); std::string content = "abcdefghijklmnopqrstuvwxyz"; - // A single 25-byte range exceeds range_size_limit, so Init() coalesces it + // A single 25-byte range exceeds range_size_limit, so AddRanges() coalesces and splits it // into three adjacent entries: {0,10}, {10,10} and {20,5}. std::shared_ptr cache_ptr = CreateTestFileAndCache("data_file", content, config, {{0, 25}}); @@ -380,7 +390,7 @@ TEST(TestReadAheadCache, TestInFlightEntryServesRacingReader) { auto gated = std::make_shared(OpenTestFile(&dir, "data_file", content)); ReadAheadCache cache(gated, config, /*file_size=*/0, TestPool()); - ASSERT_OK(cache.Init({{0, 5}})); + ASSERT_OK(cache.AddRanges({{0, 5}}).status()); cache.Warmup(); // The prefetch entry is published, but its fetch is still held. @@ -430,7 +440,7 @@ TEST(TestReadAheadCache, TestPrefetchBufferKeepsPoolAliveAfterCacheIsGone) { // Declared after the pool, so the cache is destroyed before the last // reference of this test to the pool is dropped. ReadAheadCache cache(gated, config, /*file_size=*/0, pool); - ASSERT_OK(cache.Init({{0, 5}})); + ASSERT_OK(cache.AddRanges({{0, 5}}).status()); cache.Warmup(); ASSERT_EQ(gated->AsyncReadCount(), 1); // The fetch is completed, but the stream keeps its callback - and with it @@ -474,8 +484,8 @@ TEST(TestReadAheadCache, TestPreBufferWindowLimit) { ASSERT_EQ(io_hook->IOCount(), 0); } -// Test that Init() rejects a second call until the cache is reset. -TEST(TestReadAheadCache, TestDoubleInit) { +// Test that the cache can be re-initialized after Reset() and serves the new ranges. +TEST(TestReadAheadCache, TestReinitAfterReset) { CacheConfig config = TestCacheConfig(/*range_size_limit=*/10, /*hole_size_limit=*/2, /*pre_buffer_limit=*/1024); std::string content = "abcdefghijklmnopqrstuvwxyz"; @@ -483,35 +493,220 @@ TEST(TestReadAheadCache, TestDoubleInit) { CreateTestFileAndCache("data_file", content, config, {{0, 5}}); ReadAheadCache& cache = *cache_ptr; - Status status = cache.Init({{8, 5}}); - ASSERT_FALSE(status.ok()); - - // The original ranges still work. AssertReadEquals({0, 5}, "abcde", &cache); + + cache.Reset(); + ASSERT_OK(cache.AddRanges({{3, 4}}).status()); + AssertReadEquals({3, 4}, "defg", &cache); + + // The old ranges are gone. + AssertReadMiss({20, 2}, &cache); } -// Test that the cache can be re-initialized after Reset() and serves the new ranges. -TEST(TestReadAheadCache, TestReinitAfterReset) { +// Test that AddRanges() registers ranges that only become known after the cache was +// initialized, so the reads that follow are served from the cache. +TEST(TestReadAheadCache, TestAddRangesRegistersNewRanges) { CacheConfig config = TestCacheConfig(/*range_size_limit=*/10, - /*hole_size_limit=*/2, /*pre_buffer_limit=*/1024); + /*hole_size_limit=*/0, /*pre_buffer_limit=*/1024); + std::string content = "abcdefghijklmnopqrstuvwxyz"; + std::shared_ptr cache_ptr = + CreateTestFileAndCache("data_file", content, config, {{0, 10}}); + ReadAheadCache& cache = *cache_ptr; + + AssertReadMiss({16, 10}, &cache); + + ASSERT_OK_AND_ASSIGN(std::optional first_added, cache.AddRanges({{16, 10}})); + ASSERT_TRUE(first_added.has_value()); + ASSERT_EQ(16u, first_added.value()); + + AssertReadEquals({16, 10}, "qrstuvwxyz", &cache); + // The range registered first is still served. + AssertReadEquals({0, 10}, "abcdefghij", &cache); +} + +// Of a new range intersecting an already registered one, only the part the registered ranges do +// not cover is registered: the registered ranges stay disjoint, and the read the new range asked +// for is served instead of dropped along with it. +TEST(TestReadAheadCache, TestAddRangesClipsIntersectingRange) { + CacheConfig config = TestCacheConfig(/*range_size_limit=*/10, + /*hole_size_limit=*/0, /*pre_buffer_limit=*/1024); + std::string content = "abcdefghijklmnopqrstuvwxyz"; + std::shared_ptr cache_ptr = + CreateTestFileAndCache("data_file", content, config, {{0, 10}}); + ReadAheadCache& cache = *cache_ptr; + + ASSERT_OK_AND_ASSIGN(std::optional first_added, cache.AddRanges({{5, 10}})); + ASSERT_TRUE(first_added.has_value()); + // [5, 10) is already registered, so the registration starts at the first byte that is not. + ASSERT_EQ(10u, first_added.value()); + + // Served by the registered range and the newly registered one together, as they are adjacent. + AssertReadEquals({5, 10}, "fghijklmno", &cache); + AssertReadEquals({0, 10}, "abcdefghij", &cache); +} + +// A new range fully covered by the registered ones registers nothing at all. +TEST(TestReadAheadCache, TestAddRangesDropsCoveredRange) { + CacheConfig config = TestCacheConfig(/*range_size_limit=*/10, + /*hole_size_limit=*/0, /*pre_buffer_limit=*/1024); + std::string content = "abcdefghijklmnopqrstuvwxyz"; + std::shared_ptr cache_ptr = + CreateTestFileAndCache("data_file", content, config, {{0, 10}}); + ReadAheadCache& cache = *cache_ptr; + + ASSERT_OK_AND_ASSIGN(std::optional first_added, cache.AddRanges({{2, 4}})); + ASSERT_FALSE(first_added.has_value()); + + AssertReadEquals({0, 10}, "abcdefghij", &cache); +} + +// The registered ranges stay offset-ordered and disjoint when the new ranges interleave with them, +// so that a read still finds its covering entries. +TEST(TestReadAheadCache, TestAddRangesInterleavesWithRegisteredRanges) { + CacheConfig config = TestCacheConfig(/*range_size_limit=*/26, + /*hole_size_limit=*/0, /*pre_buffer_limit=*/1024); + std::string content = "abcdefghijklmnopqrstuvwxyz"; + std::shared_ptr cache_ptr = + CreateTestFileAndCache("data_file", content, config, {{4, 4}, {16, 4}}); + ReadAheadCache& cache = *cache_ptr; + + // Two ranges each straddling a registered one, plus one entirely in the hole between them. + ASSERT_OK_AND_ASSIGN(std::optional first_added, + cache.AddRanges({{2, 4}, {10, 2}, {18, 6}})); + ASSERT_TRUE(first_added.has_value()); + ASSERT_EQ(2u, first_added.value()); + + AssertReadEquals({2, 6}, "cdefgh", &cache); + AssertReadEquals({10, 2}, "kl", &cache); + AssertReadEquals({16, 8}, "qrstuvwx", &cache); + // The holes nobody registered are still misses. + AssertReadMiss({12, 4}, &cache); +} + +// A registered range larger than the range size limit is cut into pieces, so it is fetched by +// several concurrent requests rather than by one long one. The pieces are adjacent, so a read +// spanning them is still served. +TEST(TestReadAheadCache, TestAddRangesSplitsLargeRange) { + CacheConfig config = TestCacheConfig(/*range_size_limit=*/5, + /*hole_size_limit=*/0, /*pre_buffer_limit=*/1024); std::string content = "abcdefghijklmnopqrstuvwxyz"; std::shared_ptr cache_ptr = CreateTestFileAndCache("data_file", content, config, {{0, 5}}); ReadAheadCache& cache = *cache_ptr; - AssertReadEquals({0, 5}, "abcde", &cache); + ASSERT_OK_AND_ASSIGN(std::optional first_added, cache.AddRanges({{10, 15}})); + ASSERT_TRUE(first_added.has_value()); + ASSERT_EQ(10u, first_added.value()); - cache.Reset(); - ASSERT_OK(cache.Init({{3, 4}})); - AssertReadEquals({3, 4}, "defg", &cache); + cache.Warmup(first_added.value()); + // One request per piece, instead of one request for the whole 15 bytes. + ASSERT_EQ(3u, CounterOf(ReadAheadCacheMetrics::IO_COUNT, &cache)); - // The old ranges are gone. - AssertReadMiss({20, 2}, &cache); + AssertReadEquals({10, 15}, "klmnopqrstuvwxy", &cache); + ASSERT_EQ(3u, CounterOf(ReadAheadCacheMetrics::IO_COUNT, &cache)); +} + +// AddRanges() rebuilds the per-range cached flags, so a range fetched before the call is not +// fetched a second time. +TEST(TestReadAheadCache, TestAddRangesKeepsCachedRanges) { + // A window of exactly one range, so that reading the first one never reaches the second. + CacheConfig config = TestCacheConfig(/*range_size_limit=*/10, + /*hole_size_limit=*/0, /*pre_buffer_limit=*/10); + std::string content = "abcdefghijklmnopqrstuvwxyz"; + std::shared_ptr cache_ptr = + CreateTestFileAndCache("data_file", content, config, {{0, 10}}); + ReadAheadCache& cache = *cache_ptr; + + cache.Warmup(); + AssertReadEquals({0, 10}, "abcdefghij", &cache); + + auto io_count_of = [&cache]() { return CounterOf(ReadAheadCacheMetrics::IO_COUNT, &cache); }; + const uint64_t io_count_before = io_count_of(); + ASSERT_EQ(1u, io_count_before); + + ASSERT_OK_AND_ASSIGN(std::optional first_added, cache.AddRanges({{16, 10}})); + ASSERT_TRUE(first_added.has_value()); + + AssertReadEquals({0, 10}, "abcdefghij", &cache); + ASSERT_EQ(io_count_before, io_count_of()); +} + +// Re-registering a range the cache already covers adds nothing, and must leave the published entry +// intact: the earlier fetch stays valid, so a later read of the same range is still a hit and +// issues no further IO. A no-op registration that dropped the buffer would turn this read into a +// miss (and lose the future to wait on an in-flight fetch). +TEST(TestReadAheadCache, TestReRegisteringCoveredRangeKeepsCache) { + CacheConfig config = TestCacheConfig(/*range_size_limit=*/10, + /*hole_size_limit=*/0, /*pre_buffer_limit=*/10); + std::string content = "abcdefghijklmnopqrstuvwxyz"; + std::shared_ptr cache_ptr = + CreateTestFileAndCache("data_file", content, config, {{0, 10}}); + ReadAheadCache& cache = *cache_ptr; + + cache.Warmup(); + AssertReadEquals({0, 10}, "abcdefghij", &cache); + auto io_count_of = [&cache]() { return CounterOf(ReadAheadCacheMetrics::IO_COUNT, &cache); }; + const uint64_t io_count_before = io_count_of(); + ASSERT_EQ(1u, io_count_before); + + // [2,6) is already covered by [0,10), so nothing new is registered. + ASSERT_OK_AND_ASSIGN(std::optional first_added, cache.AddRanges({{2, 6}})); + ASSERT_FALSE(first_added.has_value()); + + // The published entry survived the no-op registration: still a hit, no new IO. + AssertReadEquals({0, 10}, "abcdefghij", &cache); + ASSERT_EQ(io_count_before, io_count_of()); +} + +// An empty registration is likewise a no-op that keeps the already published entries. +TEST(TestReadAheadCache, TestEmptyAddRangesKeepsCache) { + CacheConfig config = TestCacheConfig(/*range_size_limit=*/10, + /*hole_size_limit=*/0, /*pre_buffer_limit=*/10); + std::string content = "abcdefghijklmnopqrstuvwxyz"; + std::shared_ptr cache_ptr = + CreateTestFileAndCache("data_file", content, config, {{0, 10}}); + ReadAheadCache& cache = *cache_ptr; + + cache.Warmup(); + AssertReadEquals({0, 10}, "abcdefghij", &cache); + auto io_count_of = [&cache]() { return CounterOf(ReadAheadCacheMetrics::IO_COUNT, &cache); }; + const uint64_t io_count_before = io_count_of(); + ASSERT_EQ(1u, io_count_before); + + ASSERT_OK_AND_ASSIGN(std::optional first_added, cache.AddRanges({})); + ASSERT_FALSE(first_added.has_value()); + + AssertReadEquals({0, 10}, "abcdefghij", &cache); + ASSERT_EQ(io_count_before, io_count_of()); +} + +// Warmup(offset) starts fetching from the given range rather than from the first one, which is +// what a pass registering its ranges mid-read needs: the earlier ranges belong to the pass that +// has already run. +TEST(TestReadAheadCache, TestWarmupFromOffset) { + CacheConfig config = TestCacheConfig(/*range_size_limit=*/10, + /*hole_size_limit=*/0, /*pre_buffer_limit=*/1024); + std::string content = "abcdefghijklmnopqrstuvwxyz"; + std::shared_ptr cache_ptr = + CreateTestFileAndCache("data_file", content, config, {{0, 5}, {8, 5}, {16, 5}}); + ReadAheadCache& cache = *cache_ptr; + + auto io_hook = paimon::IOHook::GetInstance(); + paimon::ScopeGuard guard([&io_hook]() { io_hook->Clear(); }); + // IOCount() only counts while armed; INT64_MAX never triggers the error mode. + io_hook->Reset(INT64_MAX, paimon::IOHook::Mode::RETURN_ERROR); + + cache.Warmup(/*from_offset=*/16); + // Only the range at 16 was fetched: the two before it were skipped. + ASSERT_EQ(io_hook->IOCount(), 1); + + AssertReadEquals({16, 5}, "qrstu", &cache); + ASSERT_EQ(io_hook->IOCount(), 1); } -// Test that Init() merges ranges separated by a small hole, so a read +// Test that AddRanges() merges ranges separated by a small hole, so a read // spanning the hole is served by the single coalesced entry. -TEST(TestReadAheadCache, TestInitCoalescesSmallHoles) { +TEST(TestReadAheadCache, TestCoalescesSmallHoles) { CacheConfig config = TestCacheConfig(/*range_size_limit=*/1024, /*hole_size_limit=*/2, /*pre_buffer_limit=*/1024); std::string content = "abcdefghijklmnopqrstuvwxyz"; @@ -639,7 +834,7 @@ TEST(TestReadAheadCache, TestBlockCacheDeclinedReadIsAMiss) { ASSERT_EQ(misses, 1u); } -// The blocks belong to the file rather than to a registration round: they +// The blocks belong to the file rather than to the registered ranges: they // survive Reset() and are only released by ReleaseBuffers(). TEST(TestReadAheadCache, TestBlockCacheSurvivesReset) { std::string content = "abcdefghijklmnopqrstuvwxyz"; diff --git a/src/paimon/core/operation/commit/retry_waiter_test.cpp b/src/paimon/core/operation/commit/retry_waiter_test.cpp deleted file mode 100644 index 2d8ea105c..000000000 --- a/src/paimon/core/operation/commit/retry_waiter_test.cpp +++ /dev/null @@ -1,50 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you 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. - */ - -#include "paimon/core/operation/commit/retry_waiter.h" - -#include - -#include "gtest/gtest.h" - -namespace paimon::test { - -TEST(RetryWaiterTest, TestRetryWaitWithZeroBoundsReturnsQuickly) { - RetryWaiter waiter(/*min_retry_wait_ms=*/0, /*max_retry_wait_ms=*/0); - - auto begin = std::chrono::steady_clock::now(); - waiter.RetryWait(/*retry_count=*/3); - auto elapsed = std::chrono::duration_cast( - std::chrono::steady_clock::now() - begin); - - ASSERT_LT(elapsed.count(), 20); -} - -TEST(RetryWaiterTest, TestRetryWaitWithLargeRetryCountIsClamped) { - RetryWaiter waiter(/*min_retry_wait_ms=*/10, /*max_retry_wait_ms=*/10); - - auto begin = std::chrono::steady_clock::now(); - waiter.RetryWait(/*retry_count=*/1024); - auto elapsed = std::chrono::duration_cast( - std::chrono::steady_clock::now() - begin); - - ASSERT_GE(elapsed.count(), 8); -} - -} // namespace paimon::test 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 70048b125..78f756fa0 100644 --- a/src/paimon/format/parquet/page_filtered_row_group_reader.cpp +++ b/src/paimon/format/parquet/page_filtered_row_group_reader.cpp @@ -20,6 +20,7 @@ #include "paimon/format/parquet/page_filtered_row_group_reader.h" #include +#include #include #include @@ -44,11 +45,15 @@ namespace paimon::parquet { namespace { -/// Ceiling on the value bytes a leaf may reserve up front from column chunk metadata alone. -/// The estimate it caps is a heuristic over footer fields, which are attacker-controlled and -/// need not describe the pages this read touches, so it must not turn into an unbounded eager -/// allocation. Past this size the builder's doubling is already amortized against a large read. -constexpr int64_t kMaxMetadataValueBytesReservation = int64_t{16} * 1024 * 1024; +/// Multiple of a column chunk's compressed size that a leaf may reserve up front in value bytes. +/// The estimate is a heuristic over footer fields, and `total_uncompressed_size` is a bare claim a +/// forged footer can inflate to make a leaf eagerly allocate gigabytes for a near-empty file. +/// `total_compressed_size` is instead checked against the file length before any read (Arrow's +/// ComputeColumnChunkRange, and GetDataPageLayout below, bound a chunk and its pages by it), so a +/// factor of it is bounded by bytes that physically exist. A leaf decompressing within the factor +/// is reserved in full; a more compressible one is capped and left to the builder's amortized +/// doubling, which is cheap against the I/O and decode of a read that large. +constexpr int64_t kMaxReservationDecompressionFactor = 8; struct DataPageLayout { int64_t column_chunk_offset; @@ -373,6 +378,8 @@ Result> PageFilteredRowGroupReader::ReadFil const int64_t reserve_values = effective_ranges.RowCount(); int64_t reserve_value_bytes = 0; const int64_t chunk_bytes = column_chunk ? column_chunk->total_uncompressed_size() : 0; + const int64_t chunk_compressed_bytes = + column_chunk ? column_chunk->total_compressed_size() : 0; const int64_t chunk_values = column_chunk ? column_chunk->num_values() : 0; if (chunk_bytes > 0 && chunk_values > 0 && reserve_values > 0) { // A heuristic, not a bound. total_uncompressed_size is uncompressed but still @@ -383,11 +390,16 @@ Result> PageFilteredRowGroupReader::ReadFil // including the pages this selection skips, so it misleads when wide values sit in // skipped pages. Either direction only costs performance — the reservation is a // hint the builder grows past when short — but they are why the result is capped - // instead of trusted. Fixed-width leaves ignore the byte count entirely. + // instead of trusted, by a factor of the file-bounded compressed size rather than the + // forgeable uncompressed one (see kMaxReservationDecompressionFactor). Fixed-width + // leaves ignore the byte count entirely. const double avg = static_cast(chunk_bytes) / static_cast(chunk_values); reserve_value_bytes = std::min( {SaturatingDoubleToInteger(avg * static_cast(reserve_values)), - chunk_bytes, kMaxMetadataValueBytesReservation}); + chunk_bytes, + SaturatingDoubleToInteger( + static_cast(chunk_compressed_bytes) * + static_cast(kMaxReservationDecompressionFactor))}); } PAIMON_RETURN_NOT_OK(ExecuteSkipReadPattern(col_idx, effective_ranges, effective_total, 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 bed06bb76..3c1875ae6 100644 --- a/src/paimon/format/parquet/parquet_file_batch_reader_test.cpp +++ b/src/paimon/format/parquet/parquet_file_batch_reader_test.cpp @@ -1827,7 +1827,7 @@ TEST_F(ParquetFileBatchReaderTest, TestPreBufferRangeFeedsReadAheadCache) { for (const auto& range : pre_buffer_ranges) { byte_ranges.emplace_back(range.first, range.second); } - ASSERT_OK(cache->Init(std::move(byte_ranges))); + ASSERT_OK(cache->AddRanges(std::move(byte_ranges)).status()); // Dispatch the prefetch immediately so every pre-buffered range is covered // before the reads below consume them. cache->Warmup(); diff --git a/src/paimon/testing/mock/mock_file_batch_reader.h b/src/paimon/testing/mock/mock_file_batch_reader.h index e017d19fd..8bc2348cf 100644 --- a/src/paimon/testing/mock/mock_file_batch_reader.h +++ b/src/paimon/testing/mock/mock_file_batch_reader.h @@ -77,6 +77,26 @@ class MockFileBatchReader : public PrefetchFileBatchReader { enable_randomize_batch_size_ = enabled; } + /// Hand out a different set of pre-buffer ranges per read schema, so that a test can tell the + /// passes of a reader that switches schemas mid-file apart: the first SetReadSchema() selects + /// index 0, the second index 1, and a pass beyond the configured sets reports no range. + void SetPreBufferRangesPerSchema( + std::vector>> ranges_by_schema) { + pre_buffer_ranges_by_schema_ = std::move(ranges_by_schema); + } + + void SetPreBufferRangeStatus(const Status& status) { + pre_buffer_range_status_ = status; + } + + Result>> PreBufferRange() override { + PAIMON_RETURN_NOT_OK(pre_buffer_range_status_); + if (schema_generation_ == 0 || schema_generation_ > pre_buffer_ranges_by_schema_.size()) { + return std::vector>{}; + } + return pre_buffer_ranges_by_schema_[schema_generation_ - 1]; + } + Status SetReadSchema(::ArrowSchema* read_schema, const std::shared_ptr& predicate, const std::optional& selection_bitmap) override { PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr arrow_schema, @@ -87,6 +107,7 @@ class MockFileBatchReader : public PrefetchFileBatchReader { // late-materialization reader moving from its probe pass to its payload pass, rely on it. current_pos_ = 0; previous_batch_first_row_num_ = std::numeric_limits::max(); + schema_generation_++; return Status::OK(); } @@ -276,6 +297,11 @@ class MockFileBatchReader : public PrefetchFileBatchReader { Status next_batch_status_; bool enable_randomize_batch_size_ = true; std::vector> read_ranges_; + // The pre-buffer ranges to report, one set per SetReadSchema() call, and how many of those + // calls have been made so far. + std::vector>> pre_buffer_ranges_by_schema_; + size_t schema_generation_ = 0; + Status pre_buffer_range_status_; int32_t warmup_count_ = 0; std::mt19937 random_engine_{std::random_device{}()}; // NOLINT(whitespace/braces) }; diff --git a/src/paimon/testing/mock/mock_file_system.h b/src/paimon/testing/mock/mock_file_system.h index 8df358271..704420801 100644 --- a/src/paimon/testing/mock/mock_file_system.h +++ b/src/paimon/testing/mock/mock_file_system.h @@ -43,8 +43,13 @@ class MockInputStream : public InputStream { Result Read(char* buffer, int64_t size, int64_t offset) override { return 0; } + /// Completes the read inline. Leaving the callback uncalled would leave the promise a + /// read-ahead cache attaches to it unresolved forever, so anything waiting for that fetch - + /// the cache releasing its buffers, or a read hitting the range - would hang. void ReadAsync(char* buffer, int64_t size, int64_t offset, - std::function&& callback) override {} + std::function&& callback) override { + callback(Status::OK()); + } Status Close() override { return Status::OK(); diff --git a/src/paimon/testing/mock/mock_format_reader_builder.h b/src/paimon/testing/mock/mock_format_reader_builder.h index 168196115..ddabe8ac4 100644 --- a/src/paimon/testing/mock/mock_format_reader_builder.h +++ b/src/paimon/testing/mock/mock_format_reader_builder.h @@ -19,6 +19,8 @@ #pragma once #include +#include +#include #include "arrow/type.h" #include "paimon/format/reader_builder.h" @@ -44,13 +46,23 @@ class MockFormatReaderBuilder : public ReaderBuilder { return this; } + /// Configure every reader this builder creates, see + /// MockFileBatchReader::SetPreBufferRangesPerSchema(). + void SetPreBufferRangesPerSchema( + std::vector>> ranges_by_schema) { + pre_buffer_ranges_by_schema_ = std::move(ranges_by_schema); + } + Result> Build( const std::shared_ptr& path) const override { - if (bitmap_) { - return std::make_unique(data_, schema_, bitmap_.value(), - read_batch_size_); + std::unique_ptr reader = + bitmap_ ? std::make_unique(data_, schema_, bitmap_.value(), + read_batch_size_) + : std::make_unique(data_, schema_, read_batch_size_); + if (!pre_buffer_ranges_by_schema_.empty()) { + reader->SetPreBufferRangesPerSchema(pre_buffer_ranges_by_schema_); } - return std::make_unique(data_, schema_, read_batch_size_); + return reader; } private: @@ -58,6 +70,7 @@ class MockFormatReaderBuilder : public ReaderBuilder { std::shared_ptr schema_; std::optional bitmap_; int32_t read_batch_size_; + std::vector>> pre_buffer_ranges_by_schema_; }; } // namespace paimon::test