From dd993eb653da089f54dbe8700957a199c115e67a Mon Sep 17 00:00:00 2001 From: "yonghao.fyh" Date: Thu, 5 Nov 2026 10:38:10 +0800 Subject: [PATCH 1/6] feat(read): decouple warmup and read-ahead caching from batch prefetch Closes apache/paimon-cpp#512 --- docs/source/user_guide/metrics.rst | 5 +- docs/source/user_guide/prefetch.rst | 20 +- include/paimon/defs.h | 4 +- include/paimon/format/read_hints.h | 2 +- include/paimon/read_context.h | 5 +- .../reader/prefetch_file_batch_reader.h | 3 +- include/paimon/utils/prefetch_cache_config.h | 1 + src/paimon/CMakeLists.txt | 3 + .../reader/data_file_reader_factory.cpp | 23 +- .../common/reader/data_file_reader_factory.h | 7 +- .../reader/data_file_reader_factory_test.cpp | 15 + .../reader/delegating_prefetch_reader.h | 7 +- src/paimon/common/reader/file_read_ahead.cpp | 377 ++++++++++++++++++ src/paimon/common/reader/file_read_ahead.h | 76 ++++ .../prefetch_file_batch_reader_impl.cpp | 363 ++--------------- .../reader/prefetch_file_batch_reader_impl.h | 40 +- .../prefetch_file_batch_reader_impl_test.cpp | 180 +++++++-- .../reader/warmup_file_batch_reader.cpp | 173 ++++++++ .../common/reader/warmup_file_batch_reader.h | 76 ++++ .../reader/warmup_file_batch_reader_test.cpp | 340 ++++++++++++++++ .../data_evolution_split_read_test.cpp | 61 +-- .../core/table/format/format_table_read.cpp | 2 +- .../parquet/parquet_file_batch_reader.cpp | 2 +- .../parquet_file_batch_reader_test.cpp | 7 +- 24 files changed, 1354 insertions(+), 438 deletions(-) create mode 100644 src/paimon/common/reader/file_read_ahead.cpp create mode 100644 src/paimon/common/reader/file_read_ahead.h create mode 100644 src/paimon/common/reader/warmup_file_batch_reader.cpp create mode 100644 src/paimon/common/reader/warmup_file_batch_reader.h create mode 100644 src/paimon/common/reader/warmup_file_batch_reader_test.cpp diff --git a/docs/source/user_guide/metrics.rst b/docs/source/user_guide/metrics.rst index 72ee2a176..230a90989 100644 --- a/docs/source/user_guide/metrics.rst +++ b/docs/source/user_guide/metrics.rst @@ -86,8 +86,9 @@ These metrics are C++-only and have no counterparts in Java Paimon. Prefetch I/O ------------ -``PrefetchIoMetrics`` describes only I/O that passes through the prefetch reader's instrumented -input streams. It is not a whole-query or whole-table I/O total. All counters accumulate for the +``PrefetchIoMetrics`` describes only I/O that passes through the instrumented +file input streams, including synchronous reads with batch prefetch disabled. +It is not a whole-query or whole-table I/O total. All counters accumulate for the reader lifetime and are retained across ``SetReadSchema()`` and cache reset. Latency uses relaxed atomic count and sum counters instead of per-I/O histograms to reduce hot-path cost. Collection is disabled by default; set ``prefetch.io-metrics.enabled`` to ``true`` in the read options to enable diff --git a/docs/source/user_guide/prefetch.rst b/docs/source/user_guide/prefetch.rst index e8a3a3480..1cfd1bf4e 100644 --- a/docs/source/user_guide/prefetch.rst +++ b/docs/source/user_guide/prefetch.rst @@ -94,14 +94,18 @@ context and always uses the default level. - ``WarmupLevel::DECODED``: also start decoding the next file, which hides the most latency but uses the most memory. -Warmup goes through the prefetch reader, so it only applies to Parquet and ORC -files that are actually read with prefetching, and enabling prefetch alone does -not guarantee that. With the adaptive strategy, a file whose first read range -holds more batches than one prefetch queue can buffer is read without -prefetching, and so without warmup; the ``prefetch.adaptive-disabled-count`` -metric counts how often this happens. Append compaction uses a prefetch queue -of three batches, so it only warms files whose first read range spans no more -than three full batches, such as files with small row groups. +Warmup and the read-ahead cache work independently of batch prefetch for +Parquet and ORC. With ``EnablePrefetch(false)``, ``RAW`` still warms the cache, +and ``DECODED`` prepares only the first batch on a background thread; later +batches are read synchronously. The same behavior applies when the format or +the adaptive strategy decides against prefetching. With batch prefetch active, +``DECODED`` starts the normal parallel decode loop instead. + +The defaults are batch prefetch off, read-ahead cache on, and ``RAW`` warmup. +``WarmupLevel::NONE`` disables early warmup while retaining byte caching during +normal reads. Disabling the cache makes ``RAW`` a no-op but still permits +``DECODED``. When the shared cache is active, Parquet disables its own +pre-buffering to avoid buffering the same byte ranges twice. ``WarmupLevel::RAW`` issues its fetches from the reading thread through the file system's asynchronous read, so it only overlaps I/O when that read is diff --git a/include/paimon/defs.h b/include/paimon/defs.h index 10cc3759d..9b3a2ce93 100644 --- a/include/paimon/defs.h +++ b/include/paimon/defs.h @@ -236,8 +236,8 @@ struct PAIMON_EXPORT Options { /// Applies to supported manifest read paths with or without caching. Default value is true. static const char SCAN_MANIFEST_ENTRY_LAZY_DECODE_ENABLED[]; - /// "prefetch.io-metrics.enabled" - Whether to collect per-I/O metrics for prefetch reads. - /// Default value is false. + /// "prefetch.io-metrics.enabled" - Whether to collect per-I/O metrics for Parquet and ORC + /// reads, with or without batch prefetch. Default value is false. static const char PREFETCH_IO_METRICS_ENABLED[]; /// "read.batch-size" - Read batch size for any file format if it supports. diff --git a/include/paimon/format/read_hints.h b/include/paimon/format/read_hints.h index d60b41320..a4fbb7dd7 100644 --- a/include/paimon/format/read_hints.h +++ b/include/paimon/format/read_hints.h @@ -28,7 +28,7 @@ namespace paimon { struct PAIMON_EXPORT ReadHints { /// Whether framework-level prefetch is enabled for this read. bool prefetch_enabled = false; - /// Whether the shared read-ahead cache is enabled for this read. + /// Whether the shared read-ahead cache is active for this read, independently of prefetch. bool read_ahead_cache_enabled = false; }; diff --git a/include/paimon/read_context.h b/include/paimon/read_context.h index 81b0f8fa2..25e9d135e 100644 --- a/include/paimon/read_context.h +++ b/include/paimon/read_context.h @@ -414,11 +414,12 @@ class PAIMON_EXPORT ReadContextBuilder { /// /// A read-ahead cache is used to prebuffer data ranges before they are needed, /// which can improve read performance by reducing redundant I/O operations. + /// For Parquet and ORC it also works when batch prefetch is disabled. /// @param enabled Whether to enable the read-ahead cache (default: true) /// @return Reference to this builder for method chaining. ReadContextBuilder& SetReadAheadCacheEnabled(bool enabled); - /// Set the cache configuration for prefetch read operations. + /// Set the read-ahead cache configuration for synchronous and prefetch reads. /// /// @param config The cache configuration to use. /// @return Reference to this builder for method chaining. @@ -434,6 +435,8 @@ class PAIMON_EXPORT ReadContextBuilder { /// @note WarmupLevel::RAW warms the read-ahead cache, so it has no effect and behaves like /// WarmupLevel::NONE when the cache is off (see SetReadAheadCacheEnabled()). /// WarmupLevel::DECODED does not depend on the cache. + /// Both levels work with EnablePrefetch(false). Without prefetch, DECODED prepares only the + /// first batch; subsequent batches are read synchronously. /// @see WarmupLevel ReadContextBuilder& SetWarmupLevel(WarmupLevel level); diff --git a/include/paimon/reader/prefetch_file_batch_reader.h b/include/paimon/reader/prefetch_file_batch_reader.h index 1969194bf..ed0ce55d5 100644 --- a/include/paimon/reader/prefetch_file_batch_reader.h +++ b/include/paimon/reader/prefetch_file_batch_reader.h @@ -48,7 +48,8 @@ class PAIMON_EXPORT PrefetchMetrics { static constexpr char CONSUMER_WAIT_LATENCY_US[] = "prefetch.consumer-wait-latency-us"; }; -/// C++-only metric names for I/O observed by the prefetch reader's instrumented input streams. +/// C++-only metric names for I/O observed by instrumented file input streams, with or without +/// batch prefetch. /// Java Paimon has no corresponding metrics. /// These metrics do not represent whole-query or whole-table I/O. class PAIMON_EXPORT PrefetchIoMetrics { diff --git a/include/paimon/utils/prefetch_cache_config.h b/include/paimon/utils/prefetch_cache_config.h index 10a82f1f8..68f5e32dd 100644 --- a/include/paimon/utils/prefetch_cache_config.h +++ b/include/paimon/utils/prefetch_cache_config.h @@ -143,6 +143,7 @@ enum class PAIMON_EXPORT WarmupLevel { RAW, /// Fetch the raw bytes and start the background decode loop as well, so decoded batches are /// ready before the file is read. Hides the most latency but uses the most memory. + /// Without batch prefetch, prepares only the first batch and then continues synchronously. DECODED, }; diff --git a/src/paimon/CMakeLists.txt b/src/paimon/CMakeLists.txt index 4db62cb79..faf35d251 100644 --- a/src/paimon/CMakeLists.txt +++ b/src/paimon/CMakeLists.txt @@ -154,6 +154,8 @@ set(PAIMON_COMMON_SRCS common/reader/data_file_reader_factory.cpp common/reader/predicate_batch_reader.cpp common/reader/prefetch_file_batch_reader_impl.cpp + common/reader/file_read_ahead.cpp + common/reader/warmup_file_batch_reader.cpp common/reader/reader_utils.cpp common/reader/blob_fallback_batch_reader.cpp common/reader/blob_view_resolving_batch_reader.cpp @@ -711,6 +713,7 @@ if(PAIMON_BUILD_TESTS) common/reader/late_materializing_file_batch_reader_test.cpp common/reader/predicate_batch_reader_test.cpp common/reader/prefetch_file_batch_reader_impl_test.cpp + common/reader/warmup_file_batch_reader_test.cpp common/reader/reader_utils_test.cpp common/reader/complete_row_kind_batch_reader_test.cpp common/reader/blob_fallback_batch_reader_test.cpp diff --git a/src/paimon/common/reader/data_file_reader_factory.cpp b/src/paimon/common/reader/data_file_reader_factory.cpp index edaac8301..a2ca4b190 100644 --- a/src/paimon/common/reader/data_file_reader_factory.cpp +++ b/src/paimon/common/reader/data_file_reader_factory.cpp @@ -22,7 +22,9 @@ #include "paimon/common/data/blob_defs.h" #include "paimon/common/reader/delegating_prefetch_reader.h" +#include "paimon/common/reader/file_read_ahead.h" #include "paimon/common/reader/prefetch_file_batch_reader_impl.h" +#include "paimon/common/reader/warmup_file_batch_reader.h" #include "paimon/format/file_format.h" #include "paimon/format/file_format_factory.h" #include "paimon/format/read_hints.h" @@ -46,6 +48,11 @@ bool FormatSupportsPrefetch(const std::string& format_identifier) { format_identifier != "vortex"; } +bool FormatSupportsSynchronousWarmup(const std::string& format_identifier) { + // Other format plugins may expose only FileBatchReader. Keep their plain read path intact. + return format_identifier == "parquet" || format_identifier == "orc"; +} + } // namespace Result> DataFileReaderFactory::CreateReaderBuilder( @@ -69,7 +76,10 @@ Result> DataFileReaderFactory::CreateReaderBuilde // off when the shared read-ahead cache is reading ahead. ReadHints read_hints; read_hints.prefetch_enabled = read_options.prefetch_enabled; - read_hints.read_ahead_cache_enabled = read_options.read_ahead_cache_enabled; + read_hints.read_ahead_cache_enabled = + read_options.read_ahead_cache_enabled && + ((read_options.prefetch_enabled && FormatSupportsPrefetch(format_identifier)) || + FormatSupportsSynchronousWarmup(format_identifier)); reader_builder->WithReadHints(read_hints); return reader_builder; } @@ -92,6 +102,17 @@ Result> DataFileReaderFactory::Open( read_options.warmup_level, pool, arrow_pool)); return std::make_unique(std::move(prefetch_reader)); } + if (FormatSupportsSynchronousWarmup(format_identifier)) { + PAIMON_ASSIGN_OR_RAISE( + FileReadAhead::OpenedReaders opened, + FileReadAhead::Open(file_path, file_size, reader_builder, file_system, + /*reader_count=*/1, /*executor=*/nullptr, + read_options.read_ahead_cache_enabled, read_options.cache_config, + read_options.prefetch_io_metrics_enabled, pool)); + return std::make_unique(std::move(opened.readers[0]), + std::move(opened.read_ahead), + read_options.warmup_level, arrow_pool); + } PAIMON_ASSIGN_OR_RAISE(std::shared_ptr input_stream, file_system->Open(FileStatus(file_path, file_size))); return reader_builder->Build(input_stream); diff --git a/src/paimon/common/reader/data_file_reader_factory.h b/src/paimon/common/reader/data_file_reader_factory.h index 1b6f8cd18..d13c7ffc9 100644 --- a/src/paimon/common/reader/data_file_reader_factory.h +++ b/src/paimon/common/reader/data_file_reader_factory.h @@ -40,7 +40,7 @@ class MemoryPool; class ReaderBuilder; /// Settings applied when a data file is opened, gathered from the read context and the table -/// options. The defaults are a plain read: no cache, no prefetch. +/// options. The defaults enable byte caching and raw warmup, with synchronous batch reads. struct DataFileReadOptions { /// Block cache the format reader may put what it reads into, or null for none. std::shared_ptr cache; @@ -52,10 +52,9 @@ struct DataFileReadOptions { /// Whether the shared read-ahead cache takes over reading ahead from the format itself. bool read_ahead_cache_enabled = true; CacheConfig cache_config; - /// Whether the prefetching reader records what its I/O costs. + /// Whether the file reader records what its I/O costs. bool prefetch_io_metrics_enabled = false; - /// How far the prefetching reader prepares a file before it is read. It only matters when - /// prefetching, because that reader is the one that can start work ahead of the read. + /// How far a reader prepares a file before it is read, independently of batch prefetch. WarmupLevel warmup_level = WarmupLevel::RAW; }; diff --git a/src/paimon/common/reader/data_file_reader_factory_test.cpp b/src/paimon/common/reader/data_file_reader_factory_test.cpp index 895baad44..387831bc6 100644 --- a/src/paimon/common/reader/data_file_reader_factory_test.cpp +++ b/src/paimon/common/reader/data_file_reader_factory_test.cpp @@ -36,6 +36,21 @@ namespace paimon::test { +TEST(DataFileReaderFactoryTest, OtherFormatsKeepTheirPlainReaderWithoutPrefetch) { + auto type = arrow::struct_({arrow::field("v", arrow::int32())}); + auto data = arrow::ipc::internal::json::ArrayFromJSON(type, "[[1],[2]]").ValueOrDie(); + MockFormatReaderBuilder builder(data, type, /*read_batch_size=*/2); + auto fs = std::make_shared(); + auto pool = GetDefaultPool(); + DataFileReadOptions options; + for (const char* format : {"custom", "avro", "blob", "lance", "mosaic", "vortex"}) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, + DataFileReaderFactory::Open(format, "dummy", 0, &builder, options, fs, + nullptr, pool, GetArrowPool(pool))); + ASSERT_NE(dynamic_cast(reader.get()), nullptr); + } +} + /// `Open()` is the one place a production read builds the prefetching reader, and a level that is /// not forwarded here fails nothing: the reader warms up however its own default says, and the only /// symptom is latency and memory a query never asked for. diff --git a/src/paimon/common/reader/delegating_prefetch_reader.h b/src/paimon/common/reader/delegating_prefetch_reader.h index b91c04187..7ed6844e0 100644 --- a/src/paimon/common/reader/delegating_prefetch_reader.h +++ b/src/paimon/common/reader/delegating_prefetch_reader.h @@ -74,11 +74,6 @@ class DelegatingPrefetchReader : public FileBatchReader { } void Warmup() override { - // GetReader() bypasses the prefetch layer entirely when prefetching is off, and a - // bypassed reader has no background thread to start. - if (!prefetch_reader_->NeedPrefetch()) { - return; - } prefetch_reader_->Warmup(); } @@ -88,7 +83,7 @@ class DelegatingPrefetchReader : public FileBatchReader { if (prefetch_reader_->NeedPrefetch()) { return prefetch_reader_.get(); } else { - return prefetch_reader_->GetFirstReader(); + return prefetch_reader_->GetSynchronousReader(); } } diff --git a/src/paimon/common/reader/file_read_ahead.cpp b/src/paimon/common/reader/file_read_ahead.cpp new file mode 100644 index 000000000..f48aafb72 --- /dev/null +++ b/src/paimon/common/reader/file_read_ahead.cpp @@ -0,0 +1,377 @@ +/* + * 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/common/reader/file_read_ahead.h" + +#include +#include +#include +#include +#include + +#include "paimon/common/executor/future.h" +#include "paimon/common/io/cache_input_stream.h" +#include "paimon/common/metrics/metrics_impl.h" +#include "paimon/format/reader_builder.h" + +namespace paimon { + +struct PrefetchIoMetricsState { + std::atomic read_requests{0}; + std::atomic read_requested_bytes{0}; + std::atomic read_physical_bytes{0}; + std::atomic read_failed{0}; + std::atomic read_latency_count{0}; + std::atomic read_latency_sum_us{0}; + std::atomic async_requests{0}; + std::atomic async_requested_bytes{0}; + std::atomic async_physical_bytes{0}; + std::atomic async_completed{0}; + std::atomic async_failed{0}; + std::atomic async_pending{0}; + std::atomic async_latency_count{0}; + std::atomic async_latency_sum_us{0}; +}; + +namespace { + +constexpr std::memory_order kMetricsMemoryOrder = std::memory_order_relaxed; + +template +auto RunOrSubmit(Executor* executor, Func&& task) -> std::future { + if (executor) { + return Via(executor, std::forward(task)); + } + std::promise promise; + auto future = promise.get_future(); + promise.set_value(task()); + return future; +} + +uint64_t ElapsedMicros(const std::chrono::steady_clock::time_point& start) { + return static_cast(std::chrono::duration_cast( + std::chrono::steady_clock::now() - start) + .count()); +} + +void RecordLatency(uint64_t latency_us, std::atomic* count, + std::atomic* sum_us) { + count->fetch_add(1, kMetricsMemoryOrder); + sum_us->fetch_add(latency_us, kMetricsMemoryOrder); +} + +class MetricsInputStream : public InputStream { + public: + MetricsInputStream(const std::shared_ptr& stream, + const std::shared_ptr& metrics) + : stream_(stream), metrics_(metrics) {} + + MetricsInputStream(std::unique_ptr&& stream, + const std::shared_ptr& metrics) + : stream_(std::move(stream)), metrics_(metrics) {} + + Status Seek(int64_t offset, SeekOrigin origin) override { + return stream_->Seek(offset, origin); + } + + Result GetPos() const override { + return stream_->GetPos(); + } + + Result Read(char* buffer, int64_t size) override { + return RecordRead([&]() { return stream_->Read(buffer, size); }, size); + } + + Result Read(char* buffer, int64_t size, int64_t offset) override { + return RecordRead([&]() { return stream_->Read(buffer, size, offset); }, size); + } + + void ReadAsync(char* buffer, int64_t size, int64_t offset, + std::function&& callback) override { + metrics_->async_requests.fetch_add(1, kMetricsMemoryOrder); + metrics_->async_requested_bytes.fetch_add(static_cast(std::max(0, size)), + kMetricsMemoryOrder); + metrics_->async_pending.fetch_add(1, kMetricsMemoryOrder); + std::shared_ptr metrics = metrics_; + const auto start = std::chrono::steady_clock::now(); + stream_->ReadAsync( + buffer, size, offset, + [metrics, size, start, callback = std::move(callback)](Status status) mutable { + metrics->async_pending.fetch_sub(1, kMetricsMemoryOrder); + if (status.ok()) { + metrics->async_completed.fetch_add(1, kMetricsMemoryOrder); + metrics->async_physical_bytes.fetch_add( + static_cast(std::max(0, size)), kMetricsMemoryOrder); + } else { + metrics->async_failed.fetch_add(1, kMetricsMemoryOrder); + } + RecordLatency(ElapsedMicros(start), &metrics->async_latency_count, + &metrics->async_latency_sum_us); + callback(status); + }); + } + + Status Close() override { + return stream_->Close(); + } + + Result GetUri() const override { + return stream_->GetUri(); + } + + Result Length() const override { + return stream_->Length(); + } + + private: + template + Result RecordRead(ReadFunction&& read, int64_t size) { + metrics_->read_requests.fetch_add(1, kMetricsMemoryOrder); + metrics_->read_requested_bytes.fetch_add(static_cast(std::max(0, size)), + kMetricsMemoryOrder); + const auto start = std::chrono::steady_clock::now(); + Result result = read(); + if (result.ok()) { + metrics_->read_physical_bytes.fetch_add( + static_cast(std::max(0, result.value())), kMetricsMemoryOrder); + } else { + metrics_->read_failed.fetch_add(1, kMetricsMemoryOrder); + } + RecordLatency(ElapsedMicros(start), &metrics_->read_latency_count, + &metrics_->read_latency_sum_us); + return result; + } + + std::shared_ptr stream_; + std::shared_ptr metrics_; +}; + +} // namespace + +Result FileReadAhead::Open( + const std::string& path, int64_t file_size, const ReaderBuilder* builder, + const std::shared_ptr& fs, uint32_t reader_count, + const std::shared_ptr& executor, bool cache_enabled, const CacheConfig& cache_config, + bool enable_io_metrics, const std::shared_ptr& pool) { + if (file_size < 0 || !builder || !fs || reader_count == 0 || (!executor && reader_count > 1)) { + return Status::Invalid("invalid file read-ahead open arguments"); + } + std::shared_ptr io_metrics; + if (enable_io_metrics) { + io_metrics = std::make_shared(); + } + // Synchronous reads need no executor and open/build inline. With an executor, opening a file + // mostly waits on remote I/O, so the streams are all opened concurrently + // instead of opening the cache's stream first and paying its round trip before the readers' + // begin. Both waves run on the read context's `executor`, which is sized to the prefetch + // parallelism, so no separate pool is spun up and torn down per file and the open/build + // concurrency follows the configured parallelism instead of a fixed default. That reuse is + // safe: neither wave's tasks wait on one another or on `executor`, and Create never runs on + // an `executor` worker, so the CollectAll calls below only ever wait on tasks that progress. + // The build has to be a second wave rather than folded into the first: each build task wraps + // its stream in the cache, and a task waiting on the cache would block a worker while the + // cache's own open is still queued behind it. Wave one only opens, so its tasks never wait on + // one another. + const uint32_t open_count = cache_enabled ? reader_count + 1 : reader_count; + std::vector>>> open_futures; + open_futures.reserve(open_count); + for (uint32_t i = 0; i < open_count; i++) { + open_futures.push_back(RunOrSubmit( + executor.get(), [&fs, &path, file_size]() -> Result> { + return fs->Open(FileStatus(path, file_size)); + })); + } + // The tasks only reference locals of this frame, which stay alive because both CollectAll + // calls below drain every future before returning. + std::vector>> opened_streams = CollectAll(open_futures); + std::vector> streams; + streams.reserve(opened_streams.size()); + for (auto& opened_stream : opened_streams) { + if (!opened_stream.ok()) { + return opened_stream.status(); + } + streams.push_back(std::move(opened_stream).value()); + } + + size_t next_stream = 0; + std::shared_ptr cache; + if (cache_enabled) { + std::shared_ptr input_stream = std::move(streams[next_stream++]); + if (io_metrics) { + input_stream = std::make_shared(input_stream, io_metrics); + } + // The file size lets the cache align its blocks to the end of the file, + // where the metadata the readers read before any range is registered + // lives. A zero size means unknown and disables the block cache. + cache = std::make_shared(input_stream, cache_config, + static_cast(file_size), pool); + } + // Wave two builds the readers, which reads each file's footer and so still needs the cache + // above to exist. The builds are concurrent, like the opens were. + std::vector>>> futures; + futures.reserve(reader_count); + for (uint32_t i = 0; i < reader_count; i++) { + futures.push_back(RunOrSubmit( + executor.get(), + [&builder, &cache, &streams, i, next_stream, + io_metrics]() -> Result> { + std::unique_ptr input_stream = std::move(streams[next_stream + i]); + if (io_metrics) { + input_stream = + std::make_unique(std::move(input_stream), io_metrics); + } + auto cache_input_stream = + std::make_shared(std::move(input_stream), cache); + return builder->Build(cache_input_stream); + })); + } + std::vector> readers; + for (auto& file_batch_reader : CollectAll(futures)) { + if (!file_batch_reader.ok()) { + return file_batch_reader.status(); + } + std::shared_ptr reader = std::move(file_batch_reader).value(); + auto prefetch_file_batch_reader = + std::dynamic_pointer_cast(reader); + if (prefetch_file_batch_reader == nullptr) { + return Status::Invalid( + "failed to cast to prefetch file batch reader. file format not support prefetch"); + } + readers.emplace_back(prefetch_file_batch_reader); + } + auto state = std::shared_ptr(new FileReadAhead(cache, io_metrics)); + for (const auto& reader : readers) { + std::weak_ptr weak_state = state; + reader->SetPreBufferRangeCallback( + [weak_state](std::vector>&& ranges) { + if (auto shared_state = weak_state.lock()) { + shared_state->RegisterLateRanges(std::move(ranges)); + } + }); + } + return OpenedReaders{std::move(readers), std::move(state)}; +} + +Status FileReadAhead::Warmup(PrefetchFileBatchReader* reader) { + std::lock_guard lock(mutex_); + if (!cache_ || warmed_) { + return status_; + } + warmed_ = true; + auto ranges = reader->PreBufferRange(); + if (!ranges.ok()) { + status_ = ranges.status(); + ++error_count_; + return status_; + } + std::vector byte_ranges; + for (const auto& range : ranges.value()) { + byte_ranges.emplace_back(range.first, range.second); + } + auto added = cache_->AddRanges(std::move(byte_ranges)); + if (!added.ok()) { + status_ = added.status(); + ++error_count_; + return status_; + } + cache_->Warmup(); + return Status::OK(); +} + +void FileReadAhead::RegisterLateRanges(std::vector>&& ranges) { + std::lock_guard lock(mutex_); + if (!cache_ || !warmed_ || ranges.empty() || !status_.ok()) { + return; + } + std::vector byte_ranges; + for (const auto& range : ranges) { + byte_ranges.emplace_back(range.first, range.second); + } + auto added = cache_->AddRanges(std::move(byte_ranges)); + if (!added.ok()) { + status_ = added.status(); + ++error_count_; + } else if (added.value()) { + cache_->Warmup(added.value().value()); + } +} + +Status FileReadAhead::GetStatus() const { + std::lock_guard lock(mutex_); + return status_; +} + +uint64_t FileReadAhead::GetErrorCount() const { + std::lock_guard lock(mutex_); + return error_count_; +} + +void FileReadAhead::EndGeneration() { + std::lock_guard lock(mutex_); + warmed_ = false; +} + +void FileReadAhead::Reset() { + std::lock_guard lock(mutex_); + warmed_ = false; + status_ = Status::OK(); + if (cache_) { + cache_->Reset(); + } +} + +void FileReadAhead::ReleaseBuffers() { + std::lock_guard lock(mutex_); + warmed_ = false; + if (cache_) { + cache_->ReleaseBuffers(); + } +} + +std::shared_ptr FileReadAhead::GetMetrics() const { + auto result = std::make_shared(); + if (cache_) { + std::shared_ptr cache_metrics = std::make_shared(); + cache_->CollectMetrics(&cache_metrics); + result->Merge(cache_metrics); + } + if (!io_metrics_) { + return result; + } + auto set_io_counter = [&result](const char* name, const std::atomic& value) { + result->SetCounter(name, value.load(kMetricsMemoryOrder)); + }; + set_io_counter(PrefetchIoMetrics::READ_REQUESTS, io_metrics_->read_requests); + set_io_counter(PrefetchIoMetrics::READ_REQUESTED_BYTES, io_metrics_->read_requested_bytes); + set_io_counter(PrefetchIoMetrics::READ_PHYSICAL_BYTES, io_metrics_->read_physical_bytes); + set_io_counter(PrefetchIoMetrics::READ_FAILED, io_metrics_->read_failed); + set_io_counter(PrefetchIoMetrics::READ_LATENCY_COUNT, io_metrics_->read_latency_count); + set_io_counter(PrefetchIoMetrics::READ_LATENCY_SUM_US, io_metrics_->read_latency_sum_us); + set_io_counter(PrefetchIoMetrics::ASYNC_REQUESTS, io_metrics_->async_requests); + set_io_counter(PrefetchIoMetrics::ASYNC_REQUESTED_BYTES, io_metrics_->async_requested_bytes); + set_io_counter(PrefetchIoMetrics::ASYNC_PHYSICAL_BYTES, io_metrics_->async_physical_bytes); + set_io_counter(PrefetchIoMetrics::ASYNC_COMPLETED, io_metrics_->async_completed); + set_io_counter(PrefetchIoMetrics::ASYNC_FAILED, io_metrics_->async_failed); + set_io_counter(PrefetchIoMetrics::ASYNC_LATENCY_COUNT, io_metrics_->async_latency_count); + set_io_counter(PrefetchIoMetrics::ASYNC_LATENCY_SUM_US, io_metrics_->async_latency_sum_us); + result->SetGauge(PrefetchIoMetrics::ASYNC_PENDING, + static_cast(io_metrics_->async_pending.load(kMetricsMemoryOrder))); + return result; +} + +} // namespace paimon diff --git a/src/paimon/common/reader/file_read_ahead.h b/src/paimon/common/reader/file_read_ahead.h new file mode 100644 index 000000000..0979b994c --- /dev/null +++ b/src/paimon/common/reader/file_read_ahead.h @@ -0,0 +1,76 @@ +/* + * 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. + */ +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "paimon/common/utils/read_ahead_cache.h" +#include "paimon/reader/prefetch_file_batch_reader.h" + +namespace paimon { + +class Executor; +class ReaderBuilder; +struct PrefetchIoMetricsState; + +/// Per-file byte caching and I/O instrumentation, shared by synchronous and parallel readers. +/// Callers must stop their reader tasks before resetting or releasing a generation. +class FileReadAhead { + public: + struct OpenedReaders { + std::vector> readers; + std::shared_ptr read_ahead; + }; + + static Result Open(const std::string& path, int64_t file_size, + const ReaderBuilder* builder, + const std::shared_ptr& fs, uint32_t reader_count, + const std::shared_ptr& executor, bool cache_enabled, + const CacheConfig& cache_config, bool enable_io_metrics, + const std::shared_ptr& pool); + + Status Warmup(PrefetchFileBatchReader* reader); + Status GetStatus() const; + uint64_t GetErrorCount() const; + void EndGeneration(); + void Reset(); + void ReleaseBuffers(); + std::shared_ptr GetMetrics() const; + + private: + FileReadAhead(const std::shared_ptr& cache, + const std::shared_ptr& io_metrics) + : cache_(cache), io_metrics_(io_metrics) {} + + void RegisterLateRanges(std::vector>&& ranges); + + std::shared_ptr cache_; + std::shared_ptr io_metrics_; + mutable std::mutex mutex_; + bool warmed_ = false; + Status status_; + uint64_t error_count_ = 0; +}; + +} // namespace paimon 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 4e5e540e4..36e712420 100644 --- a/src/paimon/common/reader/prefetch_file_batch_reader_impl.cpp +++ b/src/paimon/common/reader/prefetch_file_batch_reader_impl.cpp @@ -27,8 +27,8 @@ #include "arrow/array/array_base.h" #include "arrow/c/abi.h" #include "arrow/c/bridge.h" +#include "fmt/format.h" #include "paimon/common/executor/future.h" -#include "paimon/common/io/cache_input_stream.h" #include "paimon/common/metrics/metrics_impl.h" #include "paimon/common/reader/reader_utils.h" #include "paimon/common/utils/arrow/arrow_utils.h" @@ -62,23 +62,6 @@ struct PrefetchMetricsState { std::shared_ptr histograms = std::make_shared(); }; -struct PrefetchIoMetricsState { - std::atomic read_requests{0}; - std::atomic read_requested_bytes{0}; - std::atomic read_physical_bytes{0}; - std::atomic read_failed{0}; - std::atomic read_latency_count{0}; - std::atomic read_latency_sum_us{0}; - std::atomic async_requests{0}; - std::atomic async_requested_bytes{0}; - std::atomic async_physical_bytes{0}; - std::atomic async_completed{0}; - std::atomic async_failed{0}; - std::atomic async_pending{0}; - std::atomic async_latency_count{0}; - std::atomic async_latency_sum_us{0}; -}; - namespace { // Metrics do not synchronize reader state. A concurrently collected snapshot may be approximate. @@ -97,98 +80,6 @@ void UpdateMax(std::atomic* target, uint64_t value) { } } -void RecordLatency(uint64_t latency_us, std::atomic* count, - std::atomic* sum_us) { - count->fetch_add(1, kMetricsMemoryOrder); - sum_us->fetch_add(latency_us, kMetricsMemoryOrder); -} - -class MetricsInputStream : public InputStream { - public: - MetricsInputStream(const std::shared_ptr& stream, - const std::shared_ptr& metrics) - : stream_(stream), metrics_(metrics) {} - - MetricsInputStream(std::unique_ptr&& stream, - const std::shared_ptr& metrics) - : stream_(std::move(stream)), metrics_(metrics) {} - - Status Seek(int64_t offset, SeekOrigin origin) override { - return stream_->Seek(offset, origin); - } - - Result GetPos() const override { - return stream_->GetPos(); - } - - Result Read(char* buffer, int64_t size) override { - return RecordRead([&]() { return stream_->Read(buffer, size); }, size); - } - - Result Read(char* buffer, int64_t size, int64_t offset) override { - return RecordRead([&]() { return stream_->Read(buffer, size, offset); }, size); - } - - void ReadAsync(char* buffer, int64_t size, int64_t offset, - std::function&& callback) override { - metrics_->async_requests.fetch_add(1, kMetricsMemoryOrder); - metrics_->async_requested_bytes.fetch_add(static_cast(std::max(0, size)), - kMetricsMemoryOrder); - metrics_->async_pending.fetch_add(1, kMetricsMemoryOrder); - std::shared_ptr metrics = metrics_; - const auto start = std::chrono::steady_clock::now(); - stream_->ReadAsync( - buffer, size, offset, - [metrics, size, start, callback = std::move(callback)](Status status) mutable { - metrics->async_pending.fetch_sub(1, kMetricsMemoryOrder); - if (status.ok()) { - metrics->async_completed.fetch_add(1, kMetricsMemoryOrder); - metrics->async_physical_bytes.fetch_add( - static_cast(std::max(0, size)), kMetricsMemoryOrder); - } else { - metrics->async_failed.fetch_add(1, kMetricsMemoryOrder); - } - RecordLatency(ElapsedMicros(start), &metrics->async_latency_count, - &metrics->async_latency_sum_us); - callback(status); - }); - } - - Status Close() override { - return stream_->Close(); - } - - Result GetUri() const override { - return stream_->GetUri(); - } - - Result Length() const override { - return stream_->Length(); - } - - private: - template - Result RecordRead(ReadFunction&& read, int64_t size) { - metrics_->read_requests.fetch_add(1, kMetricsMemoryOrder); - metrics_->read_requested_bytes.fetch_add(static_cast(std::max(0, size)), - kMetricsMemoryOrder); - const auto start = std::chrono::steady_clock::now(); - Result result = read(); - if (result.ok()) { - metrics_->read_physical_bytes.fetch_add( - static_cast(std::max(0, result.value())), kMetricsMemoryOrder); - } else { - metrics_->read_failed.fetch_add(1, kMetricsMemoryOrder); - } - RecordLatency(ElapsedMicros(start), &metrics_->read_latency_count, - &metrics_->read_latency_sum_us); - return result; - } - - std::shared_ptr stream_; - std::shared_ptr metrics_; -}; - std::pair ComputeBatchSliceByReadRange( const std::vector& global_row_ids, const std::pair& read_range) { auto begin_it = @@ -230,90 +121,12 @@ Result> PrefetchFileBatchReaderImpl return Status::Invalid("executor should not be nullptr."); } - std::shared_ptr io_metrics; - if (enable_io_metrics) { - io_metrics = std::make_shared(); - } - // Opening a file mostly waits on remote I/O, so the streams are all opened concurrently - // instead of opening the cache's stream first and paying its round trip before the readers' - // begin. Both waves run on the read context's `executor`, which is sized to the prefetch - // parallelism, so no separate pool is spun up and torn down per file and the open/build - // concurrency follows the configured parallelism instead of a fixed default. That reuse is - // safe: neither wave's tasks wait on one another or on `executor`, and Create never runs on - // an `executor` worker, so the CollectAll calls below only ever wait on tasks that progress. - // The build has to be a second wave rather than folded into the first: each build task wraps - // its stream in the cache, and a task waiting on the cache would block a worker while the - // cache's own open is still queued behind it. Wave one only opens, so its tasks never wait on - // one another. - const uint32_t open_count = - read_ahead_cache_enabled ? prefetch_max_parallel_num + 1 : prefetch_max_parallel_num; - std::vector>>> open_futures; - open_futures.reserve(open_count); - for (uint32_t i = 0; i < open_count; i++) { - open_futures.push_back( - Via(executor.get(), - [&fs, &data_file_path, data_file_size]() -> Result> { - return fs->Open(FileStatus(data_file_path, data_file_size)); - })); - } - // The tasks only reference locals of this frame, which stay alive because both CollectAll - // calls below drain every future before returning. - std::vector>> opened_streams = CollectAll(open_futures); - std::vector> streams; - streams.reserve(opened_streams.size()); - for (auto& opened_stream : opened_streams) { - if (!opened_stream.ok()) { - return opened_stream.status(); - } - streams.push_back(std::move(opened_stream).value()); - } - - size_t next_stream = 0; - std::shared_ptr cache; - if (read_ahead_cache_enabled) { - std::shared_ptr input_stream = std::move(streams[next_stream++]); - if (io_metrics) { - input_stream = std::make_shared(input_stream, io_metrics); - } - // The file size lets the cache align its blocks to the end of the file, - // where the metadata the readers read before any range is registered - // lives. A zero size means unknown and disables the block cache. - cache = std::make_shared(input_stream, cache_config, - static_cast(data_file_size), pool); - } - // Wave two builds the readers, which reads each file's footer and so still needs the cache - // above to exist. The builds are concurrent, like the opens were. - std::vector>>> futures; - futures.reserve(prefetch_max_parallel_num); - for (uint32_t i = 0; i < prefetch_max_parallel_num; i++) { - futures.push_back(Via(executor.get(), - [&reader_builder, &cache, &streams, i, next_stream, - io_metrics]() -> Result> { - std::unique_ptr input_stream = - std::move(streams[next_stream + i]); - if (io_metrics) { - input_stream = std::make_unique( - std::move(input_stream), io_metrics); - } - auto cache_input_stream = std::make_shared( - std::move(input_stream), cache); - return reader_builder->Build(cache_input_stream); - })); - } - std::vector> readers; - for (auto& file_batch_reader : CollectAll(futures)) { - if (!file_batch_reader.ok()) { - return file_batch_reader.status(); - } - std::shared_ptr reader = std::move(file_batch_reader).value(); - auto prefetch_file_batch_reader = - std::dynamic_pointer_cast(reader); - if (prefetch_file_batch_reader == nullptr) { - return Status::Invalid( - "failed to cast to prefetch file batch reader. file format not support prefetch"); - } - readers.emplace_back(prefetch_file_batch_reader); - } + PAIMON_ASSIGN_OR_RAISE( + FileReadAhead::OpenedReaders opened, + FileReadAhead::Open(data_file_path, data_file_size, reader_builder, fs, + prefetch_max_parallel_num, executor, read_ahead_cache_enabled, + cache_config, enable_io_metrics, pool)); + auto& readers = opened.readers; if (prefetch_batch_count < readers.size()) { prefetch_batch_count = readers.size(); } @@ -321,17 +134,7 @@ 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)); - }); - } + opened.read_ahead, warmup_level, arrow_pool)); 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. @@ -343,19 +146,19 @@ Result> PrefetchFileBatchReaderImpl PrefetchFileBatchReaderImpl::PrefetchFileBatchReaderImpl( const std::vector>& readers, int32_t batch_size, uint32_t prefetch_queue_capacity, bool enable_adaptive_prefetch_strategy, - const std::shared_ptr& executor, const std::shared_ptr& cache, - const std::shared_ptr& io_metrics, WarmupLevel warmup_level, - const std::shared_ptr& arrow_pool) + const std::shared_ptr& executor, const std::shared_ptr& read_ahead, + WarmupLevel warmup_level, const std::shared_ptr& arrow_pool) : readers_(std::move(readers)), batch_size_(batch_size), executor_(executor), - cache_(cache), + read_ahead_(read_ahead), + synchronous_reader_(std::make_unique(readers[0], read_ahead, + warmup_level, arrow_pool)), arrow_pool_(arrow_pool), prefetch_queue_capacity_(prefetch_queue_capacity), enable_adaptive_prefetch_strategy_(enable_adaptive_prefetch_strategy), warmup_level_(warmup_level), - prefetch_metrics_(std::make_shared()), - io_metrics_(io_metrics) { + prefetch_metrics_(std::make_shared()) { for (size_t i = 0; i < readers_.size(); i++) { prefetch_queues_.emplace_back(std::make_unique>()); readers_pos_.emplace_back(std::make_unique>(0)); @@ -373,9 +176,7 @@ Status PrefetchFileBatchReaderImpl::SetReadSchema( ::ArrowSchema* read_schema, const std::shared_ptr& predicate, const std::optional& selection_bitmap) { PAIMON_RETURN_NOT_OK(CleanUp()); - if (cache_) { - cache_->Reset(); - } + read_ahead_->Reset(); PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr schema, arrow::ImportSchema(read_schema)); for (const auto& reader : readers_) { @@ -390,9 +191,7 @@ Status PrefetchFileBatchReaderImpl::SetReadSchema( Status PrefetchFileBatchReaderImpl::RefreshReadRanges() { PAIMON_RETURN_NOT_OK(CleanUp()); - if (cache_) { - cache_->Reset(); - } + read_ahead_->Reset(); return RefreshReadRangesAfterCleanUp(); } @@ -469,6 +268,7 @@ Status PrefetchFileBatchReaderImpl::SetReadRanges( read_ranges.push_back(eof_range); } read_ranges_freshed_ = true; + synchronous_reader_->Prepare(); return Status::OK(); } @@ -484,6 +284,7 @@ PrefetchFileBatchReaderImpl::DispatchReadRanges( } Status PrefetchFileBatchReaderImpl::CleanUp() { + synchronous_reader_->ResetWarmup(); auto clean_prefetch_queue = [this]() { for (auto& prefetch_queue : prefetch_queues_) { while (true) { @@ -520,11 +321,7 @@ Status PrefetchFileBatchReaderImpl::CleanUp() { read_ranges_in_group_.clear(); current_batch_global_row_ids_.clear(); read_ranges_freshed_ = false; - // SetReadSchema()/RefreshReadRanges() call cache_->Reset() right after CleanUp(), clearing the - // 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); + read_ahead_->EndGeneration(); clean_prefetch_queue(); prefetch_metrics_->queue_depth.store(0, kMetricsMemoryOrder); for (size_t i = 0; i < readers_pos_.size(); i++) { @@ -776,67 +573,14 @@ void PrefetchFileBatchReaderImpl::EnsureBackgroundThread() { } void PrefetchFileBatchReaderImpl::WarmCacheOnce() { - if (!cache_) { - return; - } - // 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; - } - auto read_ranges = readers_[0]->PreBufferRange(); - if (!read_ranges.ok()) { - SetReadStatus(read_ranges.status()); - return; - } - std::vector ranges; - for (const auto& read_range : read_ranges.value()) { - ranges.emplace_back(read_range.first, read_range.second); - } - Result> registered = cache_->AddRanges(std::move(ranges)); - if (!registered.ok()) { - SetReadStatus(registered.status()); - return; - } - // 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()); + // GetReadStatus() includes the shared cache status. The state counts each cache error once, + // even if RAW warmup is requested again on every batch. + (void)read_ahead_->Warmup(readers_[0].get()); } 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. + // CleanUp() has joined the background thread. Detach the weak state callbacks so retained + // sub-readers cannot report more byte ranges after this reader releases its generation. for (const auto& sub_reader : readers_) { sub_reader->SetPreBufferRangeCallback(nullptr); } @@ -853,6 +597,10 @@ void PrefetchFileBatchReaderImpl::Warmup() { if (!read_ranges_freshed_) { return; } + if (!need_prefetch_) { + synchronous_reader_->Warmup(); + return; + } if (warmup_level_ == WarmupLevel::RAW) { // Prefetch the raw bytes only, without starting the decoder. The first real read starts the // background thread, whose Workloop() finds the cache already warmed and skips re-Init. @@ -867,6 +615,9 @@ Result PrefetchFileBatchReaderImpl::NextBatchW if (!read_ranges_freshed_) { return Status::Invalid("prefetch reader read ranges are not initialized"); } + if (!need_prefetch_) { + return synchronous_reader_->NextBatchWithBitmap(); + } EnsureBackgroundThread(); const auto wait_start = std::chrono::steady_clock::now(); @@ -953,7 +704,7 @@ std::shared_ptr PrefetchFileBatchReaderImpl::GetReaderMetrics() const { if (need_prefetch_) { result->Merge(MetricsImpl::CollectReadMetrics(readers_)); } else if (!readers_.empty()) { - result->Merge(readers_[0]->GetReaderMetrics()); + result->Merge(synchronous_reader_->GetReaderMetrics()); } auto set_prefetch_counter = [&result](const char* name, const std::atomic& value) { @@ -966,7 +717,9 @@ std::shared_ptr PrefetchFileBatchReaderImpl::GetReaderMetrics() const { set_prefetch_counter(PrefetchMetrics::PRODUCED_BATCHES, prefetch_metrics_->produced_batches); set_prefetch_counter(PrefetchMetrics::CONSUMED_BATCHES, prefetch_metrics_->consumed_batches); set_prefetch_counter(PrefetchMetrics::DISCARDED_BATCHES, prefetch_metrics_->discarded_batches); - set_prefetch_counter(PrefetchMetrics::ERRORS, prefetch_metrics_->errors); + result->SetCounter( + PrefetchMetrics::ERRORS, + prefetch_metrics_->errors.load(kMetricsMemoryOrder) + read_ahead_->GetErrorCount()); set_prefetch_counter(PrefetchMetrics::ADAPTIVE_DISABLED_COUNT, prefetch_metrics_->adaptive_disabled_count); set_prefetch_counter(PrefetchMetrics::QUEUE_FULL_COUNT, prefetch_metrics_->queue_full_count); @@ -982,35 +735,9 @@ std::shared_ptr PrefetchFileBatchReaderImpl::GetReaderMetrics() const { PrefetchMetrics::QUEUE_DEPTH_MAX, static_cast(prefetch_metrics_->queue_depth_max.load(kMetricsMemoryOrder))); result->Merge(prefetch_metrics_->histograms); - if (cache_) { - // PR #209 owns the read-ahead cache metrics. Keep collecting its file-level - // hit/miss counters without defining another C++ metrics surface here. - std::shared_ptr cache_metrics = std::make_shared(); - cache_->CollectMetrics(&cache_metrics); - result->Merge(cache_metrics); - } - - if (!io_metrics_) { - return result; + if (need_prefetch_) { + result->Merge(read_ahead_->GetMetrics()); } - auto set_io_counter = [&result](const char* name, const std::atomic& value) { - result->SetCounter(name, value.load(kMetricsMemoryOrder)); - }; - set_io_counter(PrefetchIoMetrics::READ_REQUESTS, io_metrics_->read_requests); - set_io_counter(PrefetchIoMetrics::READ_REQUESTED_BYTES, io_metrics_->read_requested_bytes); - set_io_counter(PrefetchIoMetrics::READ_PHYSICAL_BYTES, io_metrics_->read_physical_bytes); - set_io_counter(PrefetchIoMetrics::READ_FAILED, io_metrics_->read_failed); - set_io_counter(PrefetchIoMetrics::READ_LATENCY_COUNT, io_metrics_->read_latency_count); - set_io_counter(PrefetchIoMetrics::READ_LATENCY_SUM_US, io_metrics_->read_latency_sum_us); - set_io_counter(PrefetchIoMetrics::ASYNC_REQUESTS, io_metrics_->async_requests); - set_io_counter(PrefetchIoMetrics::ASYNC_REQUESTED_BYTES, io_metrics_->async_requested_bytes); - set_io_counter(PrefetchIoMetrics::ASYNC_PHYSICAL_BYTES, io_metrics_->async_physical_bytes); - set_io_counter(PrefetchIoMetrics::ASYNC_COMPLETED, io_metrics_->async_completed); - set_io_counter(PrefetchIoMetrics::ASYNC_FAILED, io_metrics_->async_failed); - set_io_counter(PrefetchIoMetrics::ASYNC_LATENCY_COUNT, io_metrics_->async_latency_count); - set_io_counter(PrefetchIoMetrics::ASYNC_LATENCY_SUM_US, io_metrics_->async_latency_sum_us); - result->SetGauge(PrefetchIoMetrics::ASYNC_PENDING, - static_cast(io_metrics_->async_pending.load(kMetricsMemoryOrder))); return result; } @@ -1021,6 +748,9 @@ Result> PrefetchFileBatchReaderImpl::GetFileSchem Result PrefetchFileBatchReaderImpl::GetPreviousBatchFileRowId( uint64_t batch_row_id) const { + if (!need_prefetch_) { + return synchronous_reader_->GetPreviousBatchFileRowId(batch_row_id); + } if (current_batch_global_row_ids_.empty()) { return Status::Invalid( "Last batch is not read or last batch is empty, cannot get previous batch global row " @@ -1054,7 +784,7 @@ void PrefetchFileBatchReaderImpl::SetReadStatus(const Status& status) { Status PrefetchFileBatchReaderImpl::GetReadStatus() const { std::shared_lock lock(rw_mutex_); - return read_status_; + return read_status_.ok() ? read_ahead_->GetStatus() : read_status_; } Result PrefetchFileBatchReaderImpl::IsEofRange( const std::pair& read_range) const { @@ -1073,14 +803,9 @@ void PrefetchFileBatchReaderImpl::Close() { // GetReaderMetrics() after that. The cache is reset only when the reader is reused via // SetReadSchema()/RefreshReadRanges(). (void)CleanUp(); - if (cache_) { - // Free the prefetched buffers of this file right away (ConcatBatchReader keeps - // closed file readers alive until the whole scan finishes), but keep the - // counters for GetReaderMetrics(). - cache_->ReleaseBuffers(); - } - for (const auto& reader : readers_) { - reader->Close(); + synchronous_reader_->Close(); + for (size_t i = 1; i < readers_.size(); ++i) { + readers_[i]->Close(); } ClearPreBufferRangeCallbacks(); } 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 164cf6c6b..11d5ac232 100644 --- a/src/paimon/common/reader/prefetch_file_batch_reader_impl.h +++ b/src/paimon/common/reader/prefetch_file_batch_reader_impl.h @@ -35,7 +35,8 @@ #include #include "arrow/c/abi.h" -#include "paimon/common/utils/read_ahead_cache.h" +#include "paimon/common/reader/file_read_ahead.h" +#include "paimon/common/reader/warmup_file_batch_reader.h" #include "paimon/common/utils/threadsafe_queue.h" #include "paimon/reader/batch_reader.h" #include "paimon/reader/prefetch_file_batch_reader.h" @@ -58,7 +59,6 @@ class Executor; class Predicate; class Metrics; struct PrefetchMetricsState; -struct PrefetchIoMetricsState; class PrefetchFileBatchReaderImpl : public PrefetchFileBatchReader { public: @@ -109,6 +109,10 @@ class PrefetchFileBatchReaderImpl : public PrefetchFileBatchReader { return readers_[0].get(); } + FileBatchReader* GetSynchronousReader() const { + return synchronous_reader_.get(); + } + inline bool NeedPrefetch() const { return need_prefetch_; } @@ -123,34 +127,17 @@ class PrefetchFileBatchReaderImpl : public PrefetchFileBatchReader { PrefetchFileBatchReaderImpl( const std::vector>& readers, int32_t batch_size, uint32_t prefetch_queue_capacity, bool enable_adaptive_prefetch_strategy, - const std::shared_ptr& executor, const std::shared_ptr& cache, - const std::shared_ptr& io_metrics, WarmupLevel warmup_level, - const std::shared_ptr& arrow_pool); + const std::shared_ptr& executor, const std::shared_ptr& read_ahead, + WarmupLevel warmup_level, const std::shared_ptr& arrow_pool); Status CleanUp(); void Workloop(); /// Starts the background prefetch thread if it is not running yet. The first read does this /// itself; Warmup() is the same call made earlier, so the two must not diverge. void EnsureBackgroundThread(); - /// Initializes and warms the read-ahead cache at most once per read-range generation, using the - /// first reader's PreBufferRange(). Shared by Workloop() (DECODED) and Warmup() (RAW), so the - /// two never Init the cache twice. The flag only elects the caller that initializes, it does - /// not make a second caller wait: the two calls are ordered rather than concurrent, because a - /// RAW warmup runs on the reader's own thread before that same thread starts the background - /// thread that runs Workloop(). A no-op when there is no cache. Errors are recorded via - /// SetReadStatus() rather than returned: a warmup hint for a file that may never be read must - /// not fail an in-flight read. + /// Shares byte-range registration with the synchronous warmup path. 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(); @@ -196,7 +183,8 @@ class PrefetchFileBatchReaderImpl : public PrefetchFileBatchReader { std::mutex working_mutex_; std::condition_variable cv_; std::shared_ptr executor_; - std::shared_ptr cache_; + std::shared_ptr read_ahead_; + std::unique_ptr synchronous_reader_; std::shared_ptr arrow_pool_; mutable std::shared_mutex rw_mutex_; @@ -206,16 +194,10 @@ class PrefetchFileBatchReaderImpl : public PrefetchFileBatchReader { std::vector current_batch_global_row_ids_; 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, 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_; const WarmupLevel warmup_level_; int32_t parallel_num_; std::shared_ptr prefetch_metrics_; - std::shared_ptr io_metrics_; }; } // namespace paimon 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 a741077e9..b6c2172a7 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 @@ -25,7 +25,10 @@ #include "arrow/compute/api.h" #include "arrow/ipc/api.h" #include "gtest/gtest.h" +#include "paimon/common/reader/concat_batch_reader.h" +#include "paimon/common/reader/data_file_reader_factory.h" #include "paimon/common/reader/delegating_prefetch_reader.h" +#include "paimon/common/reader/late_materializing_reader_builder.h" #include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/common/utils/path_util.h" @@ -815,38 +818,42 @@ TEST_F(PrefetchFileBatchReaderImplTest, DelegatingPrefetchReaderWarmupForwardsWh ASSERT_TRUE(expected_array->Equals(array)); } -// With prefetching off the delegating layer hands the consumer readers_[0] directly, so forwarding -// the warmup would start a Workloop that decodes readers_[0] on a background thread while the -// consumer reads that same reader on its own thread. Warmup() must stay a no-op in that state and -// the bypassed read must still return every row. -TEST_F(PrefetchFileBatchReaderImplTest, DelegatingPrefetchReaderWarmupSkippedWhenPrefetchBypassed) { - auto data_array = PrepareArray(200); - int32_t batch_size = 10; - int32_t prefetch_max_parallel_num = 1; - ControlledMockFormatReaderBuilder reader_builder(data_array, data_type_, batch_size, - /*read_ranges=*/{{0, 100}}, - /*need_prefetch=*/true, - /*set_read_ranges_statuses=*/{}); - ASSERT_OK_AND_ASSIGN( - auto reader, - PrefetchFileBatchReaderImpl::Create( - /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, - prefetch_max_parallel_num, batch_size, - /*prefetch_batch_count=*/2, - /*enable_adaptive_prefetch_strategy=*/true, executor_, - /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, WarmupLevel::DECODED, pool_, GetArrowPool(pool_))); - auto* prefetch_reader = reader.get(); - ASSERT_FALSE(prefetch_reader->NeedPrefetch()); +// Without prefetch, the synchronous wrapper prepares one batch and joins that task before the +// consumer reads the format reader again. The parallel Workloop must never start in this state. +TEST_F(PrefetchFileBatchReaderImplTest, DelegatingPrefetchReaderWarmupWhenPrefetchBypassed) { + for (bool format_prefetch : {false, true}) { + auto data_array = PrepareArray(200); + int32_t batch_size = 10; + int32_t prefetch_max_parallel_num = 1; + ControlledMockFormatReaderBuilder reader_builder(data_array, data_type_, batch_size, + /*read_ranges=*/{{0, 100}}, + /*need_prefetch=*/format_prefetch, + /*set_read_ranges_statuses=*/{}); + ASSERT_OK_AND_ASSIGN( + auto reader, + PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, + prefetch_max_parallel_num, batch_size, + /*prefetch_batch_count=*/2, + /*enable_adaptive_prefetch_strategy=*/true, executor_, + /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), + /*enable_io_metrics=*/false, WarmupLevel::DECODED, pool_, GetArrowPool(pool_))); + auto* prefetch_reader = reader.get(); + ASSERT_FALSE(prefetch_reader->NeedPrefetch()); - DelegatingPrefetchReader delegating_reader(std::move(reader)); - delegating_reader.Warmup(); - ASSERT_FALSE(DecodeThreadStarted(prefetch_reader)); + DelegatingPrefetchReader delegating_reader(std::move(reader)); + delegating_reader.Warmup(); + ASSERT_FALSE(DecodeThreadStarted(prefetch_reader)); + ASSERT_NE(prefetch_reader->synchronous_reader_->thread_, nullptr); - ASSERT_OK_AND_ASSIGN(auto array, ReadResultCollector::CollectResult(&delegating_reader)); - auto expected_array = std::make_shared(data_array); - ASSERT_TRUE(expected_array->Equals(array)); - ASSERT_FALSE(DecodeThreadStarted(prefetch_reader)); + ASSERT_OK_AND_ASSIGN(auto array, ReadResultCollector::CollectResult(&delegating_reader)); + auto expected_array = std::make_shared(data_array); + ASSERT_TRUE(expected_array->Equals(array)); + ASSERT_FALSE(DecodeThreadStarted(prefetch_reader)); + delegating_reader.Close(); + delegating_reader.Warmup(); + ASSERT_NOK(delegating_reader.NextBatchWithBitmap()); + } } TEST_F(PrefetchFileBatchReaderImplTest, SetReadRanges) { @@ -939,6 +946,9 @@ TEST_F(PrefetchFileBatchReaderImplTest, WorkloopSetReadStatusWhenCacheInitFailed ASSERT_NOK_WITH_MSG(prefetch_reader->GetReadStatus(), "range size limit 4096 should be larger than hole size limit 8192"); + ASSERT_OK_AND_ASSIGN(uint64_t errors, + prefetch_reader->GetReaderMetrics()->GetCounter(PrefetchMetrics::ERRORS)); + ASSERT_EQ(errors, 1); } TEST_F(PrefetchFileBatchReaderImplTest, DoReadBatchReturnOkWhenShutdown) { @@ -1415,9 +1425,113 @@ TEST_P(PrefetchFileBatchReaderImplTest, TestRowMapping) { } } -// WarmupLevel::NONE makes Warmup() a complete no-op: it neither starts the background decode thread -// nor warms the read-ahead cache. Reading still returns every row, starting cold on the first -// NextBatch, which lazily starts the background loop. +TEST_P(PrefetchFileBatchReaderImplTest, FactoryWarmupIsIndependentOfPrefetch) { + const auto [format, cache_enabled] = GetParam(); + auto data = PrepareArray(90); + PrepareTestData(format, data, /*stripe_row_count=*/30, /*row_index_stride=*/30); + const std::string path = PathUtil::JoinPath(dir_->Str(), "file." + format); + ASSERT_OK_AND_ASSIGN(auto file_status, local_fs_->GetFileStatus(path)); + + for (bool late_materializing : {false, true}) { + for (int32_t mode = 0; mode < 3; ++mode) { + std::vector baseline_row_ids; + for (WarmupLevel level : {WarmupLevel::NONE, WarmupLevel::RAW, WarmupLevel::DECODED}) { + SCOPED_TRACE(format + " mode=" + std::to_string(mode) + + " level=" + std::to_string(static_cast(level)) + + " late_materializing=" + std::to_string(late_materializing)); + DataFileReadOptions options; + options.read_batch_size = 10; + options.prefetch_enabled = mode != 0; + // Synchronous reads must not validate or use these prefetch-only settings. + options.prefetch_max_parallel_num = mode == 0 ? 0 : 1; + options.prefetch_batch_count = mode == 0 ? 0 : (mode == 1 ? 6 : 1); + options.adaptive_prefetch_strategy = mode == 2; + options.read_ahead_cache_enabled = cache_enabled; + options.prefetch_io_metrics_enabled = true; + options.warmup_level = level; + ASSERT_OK_AND_ASSIGN( + std::unique_ptr builder, + DataFileReaderFactory::CreateReaderBuilder(format, {}, {}, options, pool_)); + if (late_materializing) { + builder = std::make_unique(std::move(builder), + GetArrowPool(pool_)); + } + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, + DataFileReaderFactory::Open( + format, path, file_status.GetLen(), builder.get(), options, + local_fs_, executor_, pool_, GetArrowPool(pool_))); + std::shared_ptr predicate; + std::optional bitmap; + if (late_materializing) { + predicate = PredicateBuilder::GreaterOrEqual(1, "f1", FieldType::BIGINT, + Literal(static_cast(35))); + bitmap.emplace(); + bitmap->AddRange(15, 65); + } + ArrowSchema schema; + ASSERT_TRUE(arrow::ExportSchema(*arrow::schema(fields_), &schema).ok()); + reader->Warmup(); + ASSERT_OK(reader->SetReadSchema(&schema, predicate, bitmap)); + reader->Warmup(); + reader->Warmup(); + ASSERT_OK_AND_ASSIGN(auto result, CollectResultAndRowIds(reader.get())); + auto expected = std::make_shared( + late_materializing ? data->Slice(35, 30) : data); + ASSERT_TRUE(expected->Equals(result.first)); + if (level == WarmupLevel::NONE) { + baseline_row_ids = result.second; + } else { + ASSERT_EQ(baseline_row_ids, result.second); + } + ASSERT_OK_AND_ASSIGN( + double enabled, reader->GetReaderMetrics()->GetGauge(PrefetchMetrics::ENABLED)); + ASSERT_EQ(enabled, mode == 1 ? 1.0 : 0.0); + if (cache_enabled) { + ASSERT_OK_AND_ASSIGN(uint64_t hits, reader->GetReaderMetrics()->GetCounter( + ReadAheadCacheMetrics::READ_HITS)); + ASSERT_GT(hits, 0); + } + reader->Close(); + reader->Warmup(); + } + } + } +} + +TEST_P(PrefetchFileBatchReaderImplTest, SynchronousMultiFileWarmupPreservesRows) { + const auto [format, cache_enabled] = GetParam(); + auto data = PrepareArray(30); + PrepareTestData(format, data, /*stripe_row_count=*/30, /*row_index_stride=*/10); + const std::string path = PathUtil::JoinPath(dir_->Str(), "file." + format); + ASSERT_OK_AND_ASSIGN(auto file_status, local_fs_->GetFileStatus(path)); + for (WarmupLevel level : {WarmupLevel::NONE, WarmupLevel::RAW, WarmupLevel::DECODED}) { + DataFileReadOptions options; + options.read_batch_size = 10; + options.read_ahead_cache_enabled = cache_enabled; + options.warmup_level = level; + ASSERT_OK_AND_ASSIGN( + std::unique_ptr builder, + DataFileReaderFactory::CreateReaderBuilder(format, {}, {}, options, pool_)); + std::vector> readers; + for (int32_t i = 0; i < 2; ++i) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, + DataFileReaderFactory::Open( + format, path, file_status.GetLen(), builder.get(), options, + local_fs_, executor_, pool_, GetArrowPool(pool_))); + ArrowSchema schema; + ASSERT_TRUE(arrow::ExportSchema(*arrow::schema(fields_), &schema).ok()); + ASSERT_OK(reader->SetReadSchema(&schema, nullptr, std::nullopt)); + readers.push_back(std::move(reader)); + } + ConcatBatchReader concat(std::move(readers), GetArrowPool(pool_)); + ASSERT_OK_AND_ASSIGN(auto result, ReadResultCollector::CollectResult(&concat)); + auto expected = std::make_shared(arrow::ArrayVector{data, data}); + ASSERT_TRUE(expected->Equals(result)); + concat.Close(); + } +} + +// NONE leaves the first read to start the decode loop and populate the cache. TEST_P(PrefetchFileBatchReaderImplTest, TestWarmupLevelNone) { auto [file_format, read_ahead_cache_enabled] = GetParam(); auto data_array = PrepareArray(90); diff --git a/src/paimon/common/reader/warmup_file_batch_reader.cpp b/src/paimon/common/reader/warmup_file_batch_reader.cpp new file mode 100644 index 000000000..f81646593 --- /dev/null +++ b/src/paimon/common/reader/warmup_file_batch_reader.cpp @@ -0,0 +1,173 @@ +/* + * 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 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/common/reader/warmup_file_batch_reader.h" + +#include + +#include "paimon/common/metrics/metrics_impl.h" +#include "paimon/common/reader/reader_utils.h" + +namespace paimon { + +WarmupFileBatchReader::WarmupFileBatchReader(std::shared_ptr reader, + std::shared_ptr read_ahead, + WarmupLevel level, + std::shared_ptr arrow_pool) + : reader_(std::move(reader)), + read_ahead_(std::move(read_ahead)), + level_(level), + arrow_pool_(std::move(arrow_pool)) {} + +WarmupFileBatchReader::~WarmupFileBatchReader() { + ResetWarmup(); + reader_->SetPreBufferRangeCallback(nullptr); + read_ahead_->ReleaseBuffers(); +} + +void WarmupFileBatchReader::Join() const { + if (thread_) { + thread_->join(); + thread_.reset(); + } +} + +void WarmupFileBatchReader::ResetWarmup() { + Join(); + if (batch_ && batch_->ok()) { + ReaderUtils::ReleaseReadBatch(std::move(batch_.value()).value().first); + } + batch_.reset(); + prepared_ = false; + started_ = false; + batch_returned_ = false; + read_status_ = Status::OK(); +} + +void WarmupFileBatchReader::Prepare() { + prepared_ = true; +} + +void WarmupFileBatchReader::Warmup() { + if (!prepared_ || closed_ || started_ || level_ == WarmupLevel::NONE) { + return; + } + if (level_ == WarmupLevel::RAW) { + (void)read_ahead_->Warmup(reader_.get()); + } else { + started_ = true; + thread_ = std::make_unique([this]() { batch_.emplace(ReadOneBatch()); }); + } +} + +Result WarmupFileBatchReader::ReadOneBatch() { + PAIMON_RETURN_NOT_OK(read_ahead_->Warmup(reader_.get())); + PAIMON_ASSIGN_OR_RAISE(ReadBatchWithBitmap batch, reader_->NextBatchWithBitmap()); + Status status = read_ahead_->GetStatus(); + if (!status.ok()) { + ReaderUtils::ReleaseReadBatch(std::move(batch.first)); + return status; + } + return batch; +} + +Result WarmupFileBatchReader::NextBatchWithBitmap() { + if (closed_) { + return Status::Invalid("cannot read a closed warmup reader"); + } + PAIMON_RETURN_NOT_OK(read_status_); + // A normal read is synchronous unless Warmup() explicitly prepared the first batch. + started_ = true; + Join(); + batch_returned_ = false; + Result result = batch_ ? std::move(batch_.value()) : ReadOneBatch(); + batch_.reset(); + if (result.ok() && !BatchReader::IsEofBatch(result.value())) { + batch_returned_ = true; + } else if (!result.ok()) { + read_status_ = result.status(); + } + return result; +} + +Result WarmupFileBatchReader::NextBatch() { + PAIMON_ASSIGN_OR_RAISE(ReadBatchWithBitmap batch, NextBatchWithBitmap()); + return ReaderUtils::ApplyBitmapToReadBatch(std::move(batch), arrow_pool_); +} + +Status WarmupFileBatchReader::SetReadSchema(::ArrowSchema* schema, + const std::shared_ptr& predicate, + const std::optional& bitmap) { + ResetWarmup(); + if (closed_) { + return Status::Invalid("cannot configure a closed warmup reader"); + } + read_ahead_->Reset(); + read_status_ = reader_->SetReadSchema(schema, predicate, bitmap); + PAIMON_RETURN_NOT_OK(read_status_); + Prepare(); + return Status::OK(); +} + +std::shared_ptr WarmupFileBatchReader::GetReaderMetrics() const { + // The underlying reader may be decoding. Do not inspect its mutable metrics until it stops. + // Joining is reserved for operations needing the format reader, so a live snapshot contains + // only the thread-safe cache and I/O counters while the first batch is being prepared. + auto metrics = std::make_shared(); + if (!thread_) { + metrics->Merge(reader_->GetReaderMetrics()); + } + metrics->Merge(read_ahead_->GetMetrics()); + metrics->SetGauge(PrefetchMetrics::ENABLED, 0.0); + metrics->SetGauge(PrefetchMetrics::PARALLELISM, 1.0); + return metrics; +} + +Result> WarmupFileBatchReader::GetFileSchema() const { + Join(); + return reader_->GetFileSchema(); +} + +Result WarmupFileBatchReader::GetPreviousBatchFileRowId(uint64_t batch_row_id) const { + if (!batch_returned_) { + return Status::Invalid("no previous batch returned by warmup reader"); + } + return reader_->GetPreviousBatchFileRowId(batch_row_id); +} + +Result WarmupFileBatchReader::GetNumberOfRows() const { + Join(); + return reader_->GetNumberOfRows(); +} + +bool WarmupFileBatchReader::SupportPreciseBitmapSelection() const { + Join(); + return reader_->SupportPreciseBitmapSelection(); +} + +void WarmupFileBatchReader::Close() { + if (closed_) { + return; + } + ResetWarmup(); + closed_ = true; + reader_->SetPreBufferRangeCallback(nullptr); + read_ahead_->ReleaseBuffers(); + reader_->Close(); +} + +} // namespace paimon diff --git a/src/paimon/common/reader/warmup_file_batch_reader.h b/src/paimon/common/reader/warmup_file_batch_reader.h new file mode 100644 index 000000000..98ac5a6d8 --- /dev/null +++ b/src/paimon/common/reader/warmup_file_batch_reader.h @@ -0,0 +1,76 @@ +/* + * 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. + */ +#pragma once + +#include +#include +#include + +#include "arrow/c/abi.h" +#include "paimon/common/reader/file_read_ahead.h" +#include "paimon/reader/file_batch_reader.h" + +namespace arrow { +class MemoryPool; +} // namespace arrow + +namespace paimon { + +/// A synchronous file reader that may prepare one batch before consumption starts. +class WarmupFileBatchReader : public FileBatchReader { + public: + WarmupFileBatchReader(std::shared_ptr reader, + std::shared_ptr read_ahead, WarmupLevel level, + std::shared_ptr arrow_pool); + ~WarmupFileBatchReader() override; + + Result NextBatch() override; + Result NextBatchWithBitmap() override; + std::shared_ptr GetReaderMetrics() const override; + Result> GetFileSchema() const override; + Status SetReadSchema(::ArrowSchema* schema, const std::shared_ptr& predicate, + const std::optional& bitmap) override; + Result GetPreviousBatchFileRowId(uint64_t batch_row_id) const override; + Result GetNumberOfRows() const override; + bool SupportPreciseBitmapSelection() const override; + void Warmup() override; + void Close() override; + + /// For a containing reader which configures the underlying format reader itself. + void Prepare(); + void ResetWarmup(); + + private: + void Join() const; + Result ReadOneBatch(); + + std::shared_ptr reader_; + std::shared_ptr read_ahead_; + const WarmupLevel level_; + std::shared_ptr arrow_pool_; + mutable std::unique_ptr thread_; + std::optional> batch_; + Status read_status_; + bool prepared_ = false; + bool started_ = false; + bool closed_ = false; + bool batch_returned_ = false; +}; + +} // namespace paimon diff --git a/src/paimon/common/reader/warmup_file_batch_reader_test.cpp b/src/paimon/common/reader/warmup_file_batch_reader_test.cpp new file mode 100644 index 000000000..4e86de443 --- /dev/null +++ b/src/paimon/common/reader/warmup_file_batch_reader_test.cpp @@ -0,0 +1,340 @@ +/* + * 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 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/common/reader/warmup_file_batch_reader.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "arrow/ipc/json_simple.h" +#include "gtest/gtest.h" +#include "paimon/common/reader/reader_utils.h" +#include "paimon/common/utils/arrow/mem_utils.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/common/utils/scope_guard.h" +#include "paimon/executor.h" +#include "paimon/format/reader_builder.h" +#include "paimon/testing/mock/mock_file_batch_reader.h" +#include "paimon/testing/mock/mock_file_system.h" +#include "paimon/testing/utils/gated_async_input_stream.h" +#include "paimon/testing/utils/read_result_collector.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { +namespace { + +class ObservedReader : public MockFileBatchReader { + public: + ObservedReader(const std::shared_ptr& data, + const std::shared_ptr& stream) + : MockFileBatchReader(data, data->type(), /*read_batch_size=*/2), stream_(stream) { + EnableRandomizeBatchSize(false); + } + + Result NextBatchWithBitmap() override { + const uint32_t count = ++read_count; + last_read_thread = std::this_thread::get_id(); + auto result = Read(); + if (count == 1) { + first_read_done.set_value(); + } + return result; + } + + Result>> PreBufferRange() override { + PAIMON_RETURN_NOT_OK(range_status); + return std::vector>{{0, 4}}; + } + + void SetPreBufferRangeCallback(PreBufferRangeCallback callback) override { + callback_ = std::move(callback); + } + + void ReportLateRanges(uint64_t offset) { + if (callback_) { + callback_({{offset, 4}}); + } + } + + std::atomic read_count{0}; + std::promise first_read_done; + std::thread::id last_read_thread; + bool read_bytes = false; + Status range_status; + + private: + Result Read() { + if (read_bytes) { + char bytes[4]; + PAIMON_ASSIGN_OR_RAISE(int64_t size, stream_->Read(bytes, 4, 0)); + if (size != 4) { + return Status::IOError("short observed read"); + } + } + return MockFileBatchReader::NextBatchWithBitmap(); + } + + std::shared_ptr stream_; + PreBufferRangeCallback callback_; +}; + +class ObservedBuilder : public ReaderBuilder { + public: + explicit ObservedBuilder(std::shared_ptr data) : data_(std::move(data)) {} + + ReaderBuilder* WithMemoryPool(const std::shared_ptr& pool) override { + return this; + } + + Result> Build( + const std::shared_ptr& stream) const override { + return std::make_unique(data_, stream); + } + + private: + std::shared_ptr data_; +}; + +class GatedFileSystem : public MockFileSystem { + public: + class Stream : public MockInputStream { + public: + explicit Stream(std::shared_ptr gated) : gated_(std::move(gated)) {} + void ReadAsync(char* buffer, int64_t size, int64_t offset, + std::function&& callback) override { + gated_->ReadAsync(buffer, size, offset, std::move(callback)); + } + + private: + std::shared_ptr gated_; + }; + + Result> Open(const std::string& path) const override { + return std::make_unique(gated); + } + + std::shared_ptr gated = + std::make_shared(std::make_shared()); +}; + +} // namespace + +class WarmupFileBatchReaderTest : public ::testing::Test { + public: + void SetUp() override { + type_ = arrow::struct_({arrow::field("v", arrow::int32())}); + data_ = arrow::ipc::internal::json::ArrayFromJSON(type_, "[[0],[1],[2],[3],[4],[5]]") + .ValueOrDie(); + ASSERT_OK_AND_ASSIGN(executor_, CreateDefaultExecutor(1)); + } + + Result> Open( + WarmupLevel level, bool cache_enabled, + const std::shared_ptr& fs = std::make_shared()) { + ObservedBuilder builder(data_); + PAIMON_ASSIGN_OR_RAISE(FileReadAhead::OpenedReaders opened, + FileReadAhead::Open("dummy", 0, &builder, fs, 1, executor_, + cache_enabled, CacheConfig(), true, pool_)); + observed_ = checked_pointer_cast(opened.readers[0]); + return std::make_unique(opened.readers[0], opened.read_ahead, level, + GetArrowPool(pool_)); + } + + Status Configure(FileBatchReader* reader) { + ArrowSchema schema; + PAIMON_RETURN_NOT_OK_FROM_ARROW( + arrow::ExportSchema(*arrow::schema(type_->fields()), &schema)); + return reader->SetReadSchema(&schema, nullptr, std::nullopt); + } + + uint64_t Counter(FileBatchReader* reader, const char* name) { + EXPECT_OK_AND_ASSIGN(uint64_t value, reader->GetReaderMetrics()->GetCounter(name)); + return value; + } + + std::shared_ptr pool_ = GetDefaultPool(); + std::shared_ptr executor_; + std::shared_ptr type_; + std::shared_ptr data_; + std::shared_ptr observed_; +}; + +TEST_F(WarmupFileBatchReaderTest, RawWarmsBytesWithoutDecodingAndIsIdempotent) { + auto fs = std::make_shared(); + ASSERT_OK_AND_ASSIGN(auto reader, Open(WarmupLevel::RAW, true, fs)); + ScopeGuard release([&fs]() { fs->gated->ReleaseAll(); }); + reader->Warmup(); + ASSERT_EQ(fs->gated->AsyncReadCount(), 0); + ASSERT_OK(Configure(reader.get())); + reader->Warmup(); + reader->Warmup(); + ASSERT_EQ(fs->gated->AsyncReadCount(), 1); + ASSERT_EQ(observed_->read_count.load(), 0); + fs->gated->ReleaseAll(); + observed_->read_bytes = true; + ASSERT_OK_AND_ASSIGN(auto batch, reader->NextBatchWithBitmap()); + ReaderUtils::ReleaseReadBatch(std::move(batch.first)); + ASSERT_EQ(Counter(reader.get(), ReadAheadCacheMetrics::READ_HITS), 1); + reader->Close(); + reader->Warmup(); + ASSERT_EQ(fs->gated->AsyncReadCount(), 1); + ASSERT_EQ(Counter(reader.get(), ReadAheadCacheMetrics::READ_HITS), 1); +} + +TEST_F(WarmupFileBatchReaderTest, NoneAndRawWithoutCacheDoNotStartWork) { + for (const auto& config : + {std::make_pair(WarmupLevel::NONE, true), std::make_pair(WarmupLevel::RAW, false)}) { + ASSERT_OK_AND_ASSIGN(auto reader, Open(config.first, config.second)); + ASSERT_OK(Configure(reader.get())); + reader->Warmup(); + ASSERT_EQ(observed_->read_count.load(), 0); + if (config.second) { + ASSERT_EQ(Counter(reader.get(), ReadAheadCacheMetrics::IO_COUNT), 0); + } + ASSERT_OK_AND_ASSIGN(auto result, ReadResultCollector::CollectResult(reader.get())); + ASSERT_TRUE(result->Equals(std::make_shared(data_))); + if (config.second) { + ASSERT_GT(Counter(reader.get(), ReadAheadCacheMetrics::IO_COUNT), 0); + } + } +} + +TEST_F(WarmupFileBatchReaderTest, DecodedPreparesOneBatchThenReadsSynchronously) { + for (bool cache_enabled : {false, true}) { + ASSERT_OK_AND_ASSIGN(auto reader, Open(WarmupLevel::DECODED, cache_enabled)); + ASSERT_OK(Configure(reader.get())); + auto done = observed_->first_read_done.get_future(); + reader->Warmup(); + reader->Warmup(); + ASSERT_EQ(done.wait_for(std::chrono::seconds(5)), std::future_status::ready); + ASSERT_EQ(observed_->read_count.load(), 1); + ASSERT_NOK(reader->GetPreviousBatchFileRowId(0)); + ASSERT_OK_AND_ASSIGN(auto first, reader->NextBatchWithBitmap()); + ASSERT_NE(observed_->last_read_thread, std::this_thread::get_id()); + ASSERT_EQ(observed_->read_count.load(), 1); + ASSERT_OK_AND_ASSIGN(uint64_t row_id, reader->GetPreviousBatchFileRowId(1)); + ASSERT_EQ(row_id, 1); + ReaderUtils::ReleaseReadBatch(std::move(first.first)); + reader->Warmup(); + ASSERT_OK_AND_ASSIGN(auto second, reader->NextBatchWithBitmap()); + ASSERT_EQ(observed_->read_count.load(), 2); + ASSERT_EQ(observed_->last_read_thread, std::this_thread::get_id()); + ASSERT_OK_AND_ASSIGN(row_id, reader->GetPreviousBatchFileRowId(0)); + ASSERT_EQ(row_id, 2); + ReaderUtils::ReleaseReadBatch(std::move(second.first)); + ASSERT_OK_AND_ASSIGN(double enabled, + reader->GetReaderMetrics()->GetGauge(PrefetchMetrics::ENABLED)); + ASSERT_EQ(enabled, 0.0); + } +} + +TEST_F(WarmupFileBatchReaderTest, SchemaResetDiscardsFirstBatchAndRearmsWarmup) { + ASSERT_OK_AND_ASSIGN(auto reader, Open(WarmupLevel::DECODED, true)); + ASSERT_OK(Configure(reader.get())); + reader->Warmup(); + ASSERT_OK(Configure(reader.get())); + ASSERT_NOK(reader->GetPreviousBatchFileRowId(0)); + reader->Warmup(); + ASSERT_OK_AND_ASSIGN(auto result, ReadResultCollector::CollectResult(reader.get())); + ASSERT_TRUE(result->Equals(std::make_shared(data_))); +} + +TEST_F(WarmupFileBatchReaderTest, WarmupFailureIsReturnedOnlyWhenConsumed) { + ASSERT_OK_AND_ASSIGN(auto reader, Open(WarmupLevel::DECODED, false)); + ASSERT_OK(Configure(reader.get())); + observed_->SetNextBatchStatus(Status::IOError("first batch failed")); + reader->Warmup(); + ASSERT_NOK_WITH_MSG(reader->NextBatchWithBitmap(), "first batch failed"); + observed_->SetNextBatchStatus(Status::OK()); + ASSERT_NOK_WITH_MSG(reader->NextBatchWithBitmap(), "first batch failed"); + ASSERT_OK(Configure(reader.get())); + ASSERT_OK_AND_ASSIGN(auto result, ReadResultCollector::CollectResult(reader.get())); + ASSERT_TRUE(result->Equals(std::make_shared(data_))); +} + +TEST_F(WarmupFileBatchReaderTest, EmptyFileAndUnconsumedWarmupCloseSafely) { + for (bool empty : {false, true}) { + if (empty) { + data_ = data_->Slice(0, 0); + } + ASSERT_OK_AND_ASSIGN(auto reader, Open(WarmupLevel::DECODED, true)); + ASSERT_OK(Configure(reader.get())); + reader->Warmup(); + if (empty) { + ASSERT_OK_AND_ASSIGN(auto eof, reader->NextBatchWithBitmap()); + ASSERT_TRUE(BatchReader::IsEofBatch(eof)); + } + reader->Close(); + reader->Close(); + reader->Warmup(); + ASSERT_NOK(reader->NextBatchWithBitmap()); + } + ASSERT_OK_AND_ASSIGN(auto reader, Open(WarmupLevel::DECODED, true)); + ASSERT_OK(Configure(reader.get())); + reader->Warmup(); + reader.reset(); +} + +TEST_F(WarmupFileBatchReaderTest, RawRangeAndAsyncFailuresAreDeferredUntilRead) { + { + ASSERT_OK_AND_ASSIGN(auto reader, Open(WarmupLevel::RAW, true)); + ASSERT_OK(Configure(reader.get())); + observed_->range_status = Status::IOError("range discovery failed"); + reader->Warmup(); + ASSERT_EQ(observed_->read_count.load(), 0); + ASSERT_NOK_WITH_MSG(reader->NextBatchWithBitmap(), "range discovery failed"); + } + auto fs = std::make_shared(); + ASSERT_OK_AND_ASSIGN(auto reader, Open(WarmupLevel::RAW, true, fs)); + ScopeGuard release([&fs]() { fs->gated->ReleaseAll(); }); + ASSERT_OK(Configure(reader.get())); + observed_->read_bytes = true; + reader->Warmup(); + ASSERT_EQ(fs->gated->AsyncReadCount(), 1); + fs->gated->FailAll(Status::IOError("async warmup failed")); + ASSERT_NOK_WITH_MSG(reader->NextBatchWithBitmap(), "async warmup failed"); +} + +TEST_F(WarmupFileBatchReaderTest, LateRangesAreRegisteredOnlyInActiveGeneration) { + ASSERT_OK_AND_ASSIGN(auto reader, Open(WarmupLevel::RAW, true)); + ASSERT_OK(Configure(reader.get())); + observed_->ReportLateRanges(16); + ASSERT_EQ(Counter(reader.get(), ReadAheadCacheMetrics::IO_COUNT), 0); + reader->Warmup(); + observed_->ReportLateRanges(16); + ASSERT_EQ(Counter(reader.get(), ReadAheadCacheMetrics::IO_COUNT), 2); + observed_->ReportLateRanges(16); + ASSERT_EQ(Counter(reader.get(), ReadAheadCacheMetrics::IO_COUNT), 2); + ASSERT_OK(Configure(reader.get())); + observed_->ReportLateRanges(32); + ASSERT_EQ(Counter(reader.get(), ReadAheadCacheMetrics::IO_COUNT), 0); + reader->Warmup(); + ASSERT_EQ(Counter(reader.get(), ReadAheadCacheMetrics::IO_COUNT), 1); + reader->Close(); + observed_->ReportLateRanges(32); + ASSERT_EQ(Counter(reader.get(), ReadAheadCacheMetrics::IO_COUNT), 1); +} + +} // namespace paimon::test diff --git a/src/paimon/core/operation/data_evolution_split_read_test.cpp b/src/paimon/core/operation/data_evolution_split_read_test.cpp index 3bd6beb0b..8d98655f7 100644 --- a/src/paimon/core/operation/data_evolution_split_read_test.cpp +++ b/src/paimon/core/operation/data_evolution_split_read_test.cpp @@ -1099,34 +1099,41 @@ TEST_F(DataEvolutionSplitReadTest, TestSequentialVectorReadReusesMainReader) { ])") .ValueOrDie()); ASSERT_OK_AND_ASSIGN(std::shared_ptr executor, CreateDefaultExecutor(2)); - for (const std::vector& projection : - {std::vector{"id", "embedding"}, - std::vector{"id", "embedding", "_ROW_ID", "_SEQUENCE_NUMBER"}, - std::vector{"id", "_ROW_ID", "_SEQUENCE_NUMBER"}, - std::vector{"_ROW_ID", "_SEQUENCE_NUMBER"}}) { - auto counting_fs = std::make_shared( - PathUtil::JoinPath(bucket_path, main->file_name)); - ReadContextBuilder builder(dir->Str()); - builder.WithFileSystem(counting_fs).EnablePrefetch(false).SetReadFieldNames(projection); - ASSERT_OK_AND_ASSIGN(std::shared_ptr read_context, builder.Finish()); - ASSERT_OK_AND_ASSIGN(std::shared_ptr context, - InternalReadContext::Create(read_context, table_schema, options)); - auto split_read = std::make_unique(path_factory, context, - GetDefaultPool(), executor); - ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, split_read->CreateReader(split)); - split_read.reset(); - ASSERT_OK_AND_ASSIGN(std::shared_ptr result, - ReadResultCollector::CollectResult(std::move(reader))); - ASSERT_NE(result, nullptr); - arrow::ArrayVector columns = {expected->field(0)}; - std::vector names = {"_VALUE_KIND"}; - for (const auto& name : projection) { - columns.push_back(expected->GetFieldByName(name)); - names.push_back(name); + for (bool cache_enabled : {false, true}) { + for (const std::vector& projection : + {std::vector{"id", "embedding"}, + std::vector{"id", "embedding", "_ROW_ID", "_SEQUENCE_NUMBER"}, + std::vector{"id", "_ROW_ID", "_SEQUENCE_NUMBER"}, + std::vector{"_ROW_ID", "_SEQUENCE_NUMBER"}}) { + auto counting_fs = std::make_shared( + PathUtil::JoinPath(bucket_path, main->file_name)); + ReadContextBuilder builder(dir->Str()); + builder.WithFileSystem(counting_fs) + .EnablePrefetch(false) + .SetReadAheadCacheEnabled(cache_enabled) + .SetReadFieldNames(projection); + ASSERT_OK_AND_ASSIGN(std::shared_ptr read_context, builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr context, + InternalReadContext::Create(read_context, table_schema, options)); + auto split_read = std::make_unique(path_factory, context, + GetDefaultPool(), executor); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, + split_read->CreateReader(split)); + split_read.reset(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, + ReadResultCollector::CollectResult(std::move(reader))); + ASSERT_NE(result, nullptr); + arrow::ArrayVector columns = {expected->field(0)}; + std::vector names = {"_VALUE_KIND"}; + for (const auto& name : projection) { + columns.push_back(expected->GetFieldByName(name)); + names.push_back(name); + } + auto projected = arrow::StructArray::Make(columns, names).ValueOrDie(); + ASSERT_TRUE(result->Equals(arrow::ChunkedArray(projected))) << result->ToString(); + // The cache has its own stream; neither stream is reopened for each VECTOR range. + ASSERT_EQ(counting_fs->OpenCount(), cache_enabled ? 2 : 1); } - auto projected = arrow::StructArray::Make(columns, names).ValueOrDie(); - ASSERT_TRUE(result->Equals(arrow::ChunkedArray(projected))) << result->ToString(); - ASSERT_EQ(counting_fs->OpenCount(), 1); } // A selection confined to the first VECTOR file must not include rows from the second. diff --git a/src/paimon/core/table/format/format_table_read.cpp b/src/paimon/core/table/format/format_table_read.cpp index 7be0e2f81..be93d8bdb 100644 --- a/src/paimon/core/table/format/format_table_read.cpp +++ b/src/paimon/core/table/format/format_table_read.cpp @@ -276,7 +276,7 @@ Result> FormatTableRead::TEST_Create( const std::optional>& projection, const std::shared_ptr& pool, const std::shared_ptr& predicate, bool enable_predicate_filter) { - // No context, so nothing asked for prefetch or a cache: a file is opened plainly. + // No context: use the default synchronous read, read-ahead cache and raw warmup settings. return CreateInternal(table, projection, pool, predicate, enable_predicate_filter, DataFileReadOptions(), /*executor=*/nullptr); } diff --git a/src/paimon/format/parquet/parquet_file_batch_reader.cpp b/src/paimon/format/parquet/parquet_file_batch_reader.cpp index 50f3b9735..e3311f423 100644 --- a/src/paimon/format/parquet/parquet_file_batch_reader.cpp +++ b/src/paimon/format/parquet/parquet_file_batch_reader.cpp @@ -176,7 +176,7 @@ bool IsChunkFullyDictionaryEncoded(const ::parquet::ColumnChunkMetaData& chunk) // the same byte ranges are not buffered twice. Without hints, fall back to the option. Result ResolvePreBufferEnabled(const std::map& options, const std::optional& hints) { - if (hints.has_value() && hints->prefetch_enabled && hints->read_ahead_cache_enabled) { + if (hints.has_value() && hints->read_ahead_cache_enabled) { return false; } return OptionsUtils::GetValueFromMap(options, PARQUET_READ_ENABLE_PRE_BUFFER, true); 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 703541d19..852b07027 100644 --- a/src/paimon/format/parquet/parquet_file_batch_reader_test.cpp +++ b/src/paimon/format/parquet/parquet_file_batch_reader_test.cpp @@ -1239,7 +1239,7 @@ TEST_F(ParquetFileBatchReaderTest, TestPreBufferReadHints) { pool_, options, batch_size, hints)); ASSERT_EQ(arrow_props.pre_buffer(), true); } - // Only cache enabled (prefetch disabled): fall back to the option, which defaults to true. + // The shared cache also replaces parquet pre-buffering without batch prefetch. { std::map options; ReadHints hints; @@ -1247,7 +1247,10 @@ TEST_F(ParquetFileBatchReaderTest, TestPreBufferReadHints) { hints.read_ahead_cache_enabled = true; ASSERT_OK_AND_ASSIGN(auto arrow_props, ParquetFileBatchReader::CreateArrowReaderProperties( pool_, options, batch_size, hints)); - ASSERT_EQ(arrow_props.pre_buffer(), true); + ASSERT_EQ(arrow_props.pre_buffer(), false); + ASSERT_OK_AND_ASSIGN(auto reader_props, + ParquetFileBatchReader::CreateReaderProperties(pool_, options, hints)); + ASSERT_EQ(reader_props.is_buffered_stream_enabled(), false); } // Neither active and the option explicitly disabled: honor the option. { From 4ddc5089936b46c08b48a98125dbec752e4e90df Mon Sep 17 00:00:00 2001 From: "yonghao.fyh" Date: Thu, 5 Nov 2026 11:14:49 +0800 Subject: [PATCH 2/6] refactor(read)!: remove decoded warmup mode BREAKING CHANGE: WarmupLevel::DECODED is removed. Use RAW for byte warmup or NONE to disable it. --- docs/source/user_guide/prefetch.rst | 23 ++- include/paimon/read_context.h | 11 +- include/paimon/reader/file_batch_reader.h | 9 +- include/paimon/utils/prefetch_cache_config.h | 19 +-- .../apply_bitmap_index_batch_reader_test.cpp | 3 +- .../reader/data_file_reader_factory.cpp | 6 +- .../reader/data_file_reader_factory_test.cpp | 2 +- ...e_materializing_file_batch_reader_test.cpp | 8 +- .../prefetch_file_batch_reader_impl.cpp | 22 +-- .../reader/prefetch_file_batch_reader_impl.h | 3 +- .../prefetch_file_batch_reader_impl_test.cpp | 154 +++++++----------- .../reader/warmup_file_batch_reader.cpp | 46 +----- .../common/reader/warmup_file_batch_reader.h | 9 +- .../reader/warmup_file_batch_reader_test.cpp | 93 +++++------ ...pply_deletion_vector_batch_reader_test.cpp | 2 +- .../core/operation/read_context_test.cpp | 6 +- .../core/table/system/system_table_test.cpp | 6 +- 17 files changed, 161 insertions(+), 261 deletions(-) diff --git a/docs/source/user_guide/prefetch.rst b/docs/source/user_guide/prefetch.rst index 1cfd1bf4e..d9595b09d 100644 --- a/docs/source/user_guide/prefetch.rst +++ b/docs/source/user_guide/prefetch.rst @@ -83,7 +83,7 @@ the default ``loser-tree`` sort engine additionally warms the first file of every sorted run before merging. It applies to append-table reads and append compaction as well, but does not cross split boundaries. -How far the next file is prepared is set with +Whether the next file's raw bytes are fetched early is set with ``ReadContextBuilder::SetWarmupLevel()``. Append compaction builds its own read context and always uses the default level. @@ -91,29 +91,26 @@ context and always uses the default level. - ``WarmupLevel::RAW`` (default): fetch the next file's raw bytes through the read-ahead cache without decoding them. It has no effect when the read-ahead cache is disabled with ``ReadContextBuilder::SetReadAheadCacheEnabled()``. -- ``WarmupLevel::DECODED``: also start decoding the next file, which hides the - most latency but uses the most memory. Warmup and the read-ahead cache work independently of batch prefetch for -Parquet and ORC. With ``EnablePrefetch(false)``, ``RAW`` still warms the cache, -and ``DECODED`` prepares only the first batch on a background thread; later -batches are read synchronously. The same behavior applies when the format or -the adaptive strategy decides against prefetching. With batch prefetch active, -``DECODED`` starts the normal parallel decode loop instead. +Parquet and ORC. With ``EnablePrefetch(false)``, ``RAW`` still warms the cache +and batches are read synchronously. The same behavior applies when the format +or the adaptive strategy decides against prefetching. Warmup only fetches raw +bytes; it never decodes batches. When batch prefetch is active, the first +actual read starts the parallel decode loop. The defaults are batch prefetch off, read-ahead cache on, and ``RAW`` warmup. ``WarmupLevel::NONE`` disables early warmup while retaining byte caching during -normal reads. Disabling the cache makes ``RAW`` a no-op but still permits -``DECODED``. When the shared cache is active, Parquet disables its own -pre-buffering to avoid buffering the same byte ranges twice. +normal reads. Disabling the cache makes ``RAW`` a no-op. When the shared cache +is active, Parquet disables its own pre-buffering to avoid buffering the same +byte ranges twice. ``WarmupLevel::RAW`` issues its fetches from the reading thread through the file system's asynchronous read, so it only overlaps I/O when that read is really asynchronous. The local file system serves it synchronously, so there ``RAW`` warmup reads the next file on the reading thread before the current batch is returned, adding latency instead of hiding it. On such file systems, -prefer ``WarmupLevel::DECODED``, which fetches on the background decode thread, -or ``WarmupLevel::NONE``. +use ``WarmupLevel::NONE`` to disable early fetching. Warmup preserves the returned rows and their order, but it costs memory: under ``WarmupLevel::RAW`` the next file may hold up to the read-ahead cache's diff --git a/include/paimon/read_context.h b/include/paimon/read_context.h index 25e9d135e..20dab5902 100644 --- a/include/paimon/read_context.h +++ b/include/paimon/read_context.h @@ -425,18 +425,15 @@ class PAIMON_EXPORT ReadContextBuilder { /// @return Reference to this builder for method chaining. ReadContextBuilder& WithCacheConfig(const CacheConfig& config); - /// Set how far the reader prepares the next file before it is read. + /// Set whether the reader fetches the next file's raw bytes before it is read. /// - /// Warmup overlaps remote-storage latency with the read of the current file. Higher levels hide - /// more latency but use more memory, and may warm files that a query never reads (for example - /// when a LIMIT stops the scan early). + /// Warmup can overlap remote-storage latency with reading the current file. It may fetch bytes + /// for files that a query never reads, for example when a LIMIT stops the scan early. /// @param level The warmup level to use (default: WarmupLevel::RAW). /// @return Reference to this builder for method chaining. /// @note WarmupLevel::RAW warms the read-ahead cache, so it has no effect and behaves like /// WarmupLevel::NONE when the cache is off (see SetReadAheadCacheEnabled()). - /// WarmupLevel::DECODED does not depend on the cache. - /// Both levels work with EnablePrefetch(false). Without prefetch, DECODED prepares only the - /// first batch; subsequent batches are read synchronously. + /// RAW also works with EnablePrefetch(false) and never starts batch decoding. /// @see WarmupLevel ReadContextBuilder& SetWarmupLevel(WarmupLevel level); diff --git a/include/paimon/reader/file_batch_reader.h b/include/paimon/reader/file_batch_reader.h index 32f2508c2..a042d7d72 100644 --- a/include/paimon/reader/file_batch_reader.h +++ b/include/paimon/reader/file_batch_reader.h @@ -64,9 +64,8 @@ class PAIMON_EXPORT FileBatchReader : public BatchReader { /// Get whether or not support read precisely while bitmap pushed down. virtual bool SupportPreciseBitmapSelection() const = 0; - /// Starts whatever background work this reader would otherwise start on its first read, so a - /// caller that knows this reader is next can pay that startup while still consuming the - /// previous one. + /// Fetches this file's raw bytes ahead of its first read, so a caller that knows this reader + /// is next can overlap the fetch with consuming the previous file. /// /// This is an optional hint and never changes what the reader returns: ordering, filtering and /// metrics are the same with or without it. A reader with nothing to start, or one that is not @@ -78,8 +77,8 @@ class PAIMON_EXPORT FileBatchReader : public BatchReader { /// that cannot start its work leaves it to be started by the first read, which reports the /// failure itself. /// - /// What an implementation starts is up to the implementation - a reader that prefetches on a - /// background thread starts that thread, one that reads synchronously may do nothing at all. + /// A reader may fetch through its read-ahead cache using the filesystem's asynchronous I/O. + /// Warmup does not start batch decoding. A reader without byte-warmup support does nothing. /// @warning Whatever it starts, the call itself is not thread-safe: make it from the same /// thread that reads this reader, and not concurrently with any other call on it. virtual void Warmup() {} diff --git a/include/paimon/utils/prefetch_cache_config.h b/include/paimon/utils/prefetch_cache_config.h index 68f5e32dd..e06cc129f 100644 --- a/include/paimon/utils/prefetch_cache_config.h +++ b/include/paimon/utils/prefetch_cache_config.h @@ -124,27 +124,18 @@ class PAIMON_EXPORT CacheConfig { uint64_t block_cache_limit_ = 1024 * 1024; }; -/// Controls how far a reader prepares the next file before that file is actually read. +/// Controls whether a reader fetches the next file's raw bytes before that file is read. /// -/// Warmup overlaps remote-storage latency with the read of the current file. Each level takes the -/// next file one step further along the read pipeline: `RAW` covers the remote fetch, `DECODED` -/// adds the decode on top of it. Higher levels hide more latency, but commit more memory and -/// background I/O to files that a query may end up never reading (for example when a LIMIT stops -/// the scan early). Callers can trade latency against memory by picking a level. +/// Warmup can overlap remote-storage latency with reading the current file. It may fetch bytes +/// for a file a query never consumes, for example when a LIMIT stops the scan early. enum class PAIMON_EXPORT WarmupLevel { /// Do not warm up. The next file's I/O starts only when it is actually read. This is the /// behavior from before warmup existed and uses no extra memory or background threads. NONE, /// Fetch only the next file's raw, still-compressed bytes into memory, and leave the decoder - /// alone. Overlaps the remote fetch while keeping memory lower than `DECODED`, because no - /// decoded batches are materialized ahead of the read. It fetches through the read-ahead - /// cache, so it falls back to `NONE` when that cache is disabled. This is the default, as it - /// hides the remote fetch without committing memory to decoded batches. + /// alone. No decoded batches are materialized ahead of the read. It uses the read-ahead cache, + /// so it falls back to `NONE` when that cache is disabled. This is the default level. RAW, - /// Fetch the raw bytes and start the background decode loop as well, so decoded batches are - /// ready before the file is read. Hides the most latency but uses the most memory. - /// Without batch prefetch, prepares only the first batch and then continues synchronously. - DECODED, }; } // namespace paimon diff --git a/src/paimon/common/file_index/bitmap/apply_bitmap_index_batch_reader_test.cpp b/src/paimon/common/file_index/bitmap/apply_bitmap_index_batch_reader_test.cpp index 85e3d6002..1b2be9dfa 100644 --- a/src/paimon/common/file_index/bitmap/apply_bitmap_index_batch_reader_test.cpp +++ b/src/paimon/common/file_index/bitmap/apply_bitmap_index_batch_reader_test.cpp @@ -128,8 +128,7 @@ class ApplyBitmapIndexBatchReaderTest : public ::testing::Test, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, WarmupLevel::DECODED, pool_, - GetArrowPool(pool_))); + /*enable_io_metrics=*/false, WarmupLevel::RAW, pool_, GetArrowPool(pool_))); } else { file_batch_reader = std::make_unique(data, target_type_, batch_size); diff --git a/src/paimon/common/reader/data_file_reader_factory.cpp b/src/paimon/common/reader/data_file_reader_factory.cpp index a2ca4b190..c53ab943a 100644 --- a/src/paimon/common/reader/data_file_reader_factory.cpp +++ b/src/paimon/common/reader/data_file_reader_factory.cpp @@ -48,7 +48,7 @@ bool FormatSupportsPrefetch(const std::string& format_identifier) { format_identifier != "vortex"; } -bool FormatSupportsSynchronousWarmup(const std::string& format_identifier) { +bool FormatSupportsReadAhead(const std::string& format_identifier) { // Other format plugins may expose only FileBatchReader. Keep their plain read path intact. return format_identifier == "parquet" || format_identifier == "orc"; } @@ -79,7 +79,7 @@ Result> DataFileReaderFactory::CreateReaderBuilde read_hints.read_ahead_cache_enabled = read_options.read_ahead_cache_enabled && ((read_options.prefetch_enabled && FormatSupportsPrefetch(format_identifier)) || - FormatSupportsSynchronousWarmup(format_identifier)); + FormatSupportsReadAhead(format_identifier)); reader_builder->WithReadHints(read_hints); return reader_builder; } @@ -102,7 +102,7 @@ Result> DataFileReaderFactory::Open( read_options.warmup_level, pool, arrow_pool)); return std::make_unique(std::move(prefetch_reader)); } - if (FormatSupportsSynchronousWarmup(format_identifier)) { + if (FormatSupportsReadAhead(format_identifier)) { PAIMON_ASSIGN_OR_RAISE( FileReadAhead::OpenedReaders opened, FileReadAhead::Open(file_path, file_size, reader_builder, file_system, diff --git a/src/paimon/common/reader/data_file_reader_factory_test.cpp b/src/paimon/common/reader/data_file_reader_factory_test.cpp index 387831bc6..65221bd42 100644 --- a/src/paimon/common/reader/data_file_reader_factory_test.cpp +++ b/src/paimon/common/reader/data_file_reader_factory_test.cpp @@ -68,7 +68,7 @@ TEST(DataFileReaderFactoryTest, OpenForwardsWarmupLevelToPrefetchReader) { data_type->fields()) .ValueOrDie(); - for (WarmupLevel level : {WarmupLevel::NONE, WarmupLevel::RAW, WarmupLevel::DECODED}) { + for (WarmupLevel level : {WarmupLevel::NONE, WarmupLevel::RAW}) { MockFormatReaderBuilder reader_builder(data, data_type, /*read_batch_size=*/10); DataFileReadOptions read_options; read_options.read_batch_size = 10; 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 a95fd91a9..b4295c914 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 @@ -706,7 +706,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, WorksAsInnerOfPrefetchReader) { /*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=*/false, CacheConfig(), - /*enable_io_metrics=*/false, WarmupLevel::DECODED, pool_, GetArrowPool(pool_))); + /*enable_io_metrics=*/false, WarmupLevel::RAW, pool_, GetArrowPool(pool_))); auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, "k", FieldType::BIGINT, Literal(4l)); ::ArrowSchema c_schema; @@ -742,7 +742,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, PrefetchInnerReentrantSetReadSchema /*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=*/false, CacheConfig(), - /*enable_io_metrics=*/false, WarmupLevel::DECODED, pool_, GetArrowPool(pool_))); + /*enable_io_metrics=*/false, WarmupLevel::RAW, pool_, GetArrowPool(pool_))); auto full_schema = arrow::schema(full_fields_); auto predicate1 = @@ -805,7 +805,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, PrefetchInnerParallelReadersWithSee /*prefetch_max_parallel_num=*/3, /*batch_size=*/3, /*prefetch_batch_count=*/6, /*enable_adaptive_prefetch_strategy=*/false, executor, /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/false, CacheConfig(), - /*enable_io_metrics=*/false, WarmupLevel::DECODED, pool_, GetArrowPool(pool_))); + /*enable_io_metrics=*/false, WarmupLevel::RAW, pool_, GetArrowPool(pool_))); auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, "k", FieldType::BIGINT, Literal(5l)); ::ArrowSchema c_schema; @@ -960,7 +960,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, PrefetchInnerRegistersPayloadPreBuf /*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_))); + /*enable_io_metrics=*/false, WarmupLevel::RAW, pool_, GetArrowPool(pool_))); auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, "k", FieldType::BIGINT, Literal(4l)); ::ArrowSchema c_schema; 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 36e712420..9e7ffa0bb 100644 --- a/src/paimon/common/reader/prefetch_file_batch_reader_impl.cpp +++ b/src/paimon/common/reader/prefetch_file_batch_reader_impl.cpp @@ -284,7 +284,7 @@ PrefetchFileBatchReaderImpl::DispatchReadRanges( } Status PrefetchFileBatchReaderImpl::CleanUp() { - synchronous_reader_->ResetWarmup(); + synchronous_reader_->ResetReadState(); auto clean_prefetch_queue = [this]() { for (auto& prefetch_queue : prefetch_queues_) { while (true) { @@ -336,9 +336,8 @@ Status PrefetchFileBatchReaderImpl::CleanUp() { void PrefetchFileBatchReaderImpl::Workloop() { std::vector> futures; futures.resize(readers_.size()); - // Warm the read-ahead cache before decoding. At DECODED this is the first warmup; at RAW - // Warmup() already ran it on the caller's thread and the guard makes this a no-op, so the - // cache is never Init'd twice. + // Register the file's byte ranges before decoding. An earlier RAW Warmup() may have already + // fetched them, in which case the shared state makes this a no-op. WarmCacheOnce(); while (true) { @@ -587,7 +586,7 @@ void PrefetchFileBatchReaderImpl::ClearPreBufferRangeCallbacks() { } void PrefetchFileBatchReaderImpl::Warmup() { - if (warmup_level_ == WarmupLevel::NONE) { + if (warmup_level_ != WarmupLevel::RAW) { return; } // Ranges that are not set mean the reader is either not configured yet or already cleaned up, @@ -597,18 +596,7 @@ void PrefetchFileBatchReaderImpl::Warmup() { if (!read_ranges_freshed_) { return; } - if (!need_prefetch_) { - synchronous_reader_->Warmup(); - return; - } - if (warmup_level_ == WarmupLevel::RAW) { - // Prefetch the raw bytes only, without starting the decoder. The first real read starts the - // background thread, whose Workloop() finds the cache already warmed and skips re-Init. - WarmCacheOnce(); - return; - } - // DECODED: start the background decode loop now so decoded batches are ready before the read. - EnsureBackgroundThread(); + WarmCacheOnce(); } Result PrefetchFileBatchReaderImpl::NextBatchWithBitmap() { 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 11d5ac232..b1e4a7af2 100644 --- a/src/paimon/common/reader/prefetch_file_batch_reader_impl.h +++ b/src/paimon/common/reader/prefetch_file_batch_reader_impl.h @@ -132,8 +132,7 @@ class PrefetchFileBatchReaderImpl : public PrefetchFileBatchReader { Status CleanUp(); void Workloop(); - /// Starts the background prefetch thread if it is not running yet. The first read does this - /// itself; Warmup() is the same call made earlier, so the two must not diverge. + /// Starts the background batch prefetch thread lazily on the first read. void EnsureBackgroundThread(); /// Shares byte-range registration with the synchronous warmup path. void WarmCacheOnce(); 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 b6c2172a7..43569c8d7 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 @@ -362,13 +362,7 @@ CollectResultAndRowIds(FileBatchReader* reader) { return std::make_pair(result_array, row_ids); } -/// Whether Warmup() started the background decode loop. -/// -/// The levels differ only in what Warmup() itself starts, and metrics cannot show that: the first -/// NextBatch starts the same loop lazily, so once a read is under way every level reports produced -/// batches, while right after Warmup() a DECODED warmup may not have finished its first batch yet. -/// Whether the thread exists is settled the moment Warmup() returns, so it separates the levels -/// without racing the background work. +/// Warmup must leave batch decoding to the first read, regardless of the byte-cache setting. bool DecodeThreadStarted(const PrefetchFileBatchReaderImpl* reader) { return reader->background_thread_ != nullptr; } @@ -398,7 +392,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestSimple) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, WarmupLevel::DECODED, pool_, GetArrowPool(pool_))); + /*enable_io_metrics=*/false, WarmupLevel::RAW, pool_, GetArrowPool(pool_))); if (prefetch_max_parallel_num == 1) { ASSERT_NOK( reader->GetReaderMetrics()->GetCounter(PrefetchIoMetrics::READ_LATENCY_COUNT)); @@ -426,7 +420,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestReadWithLimits) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/true, WarmupLevel::DECODED, pool_, GetArrowPool(pool_))); + /*enable_io_metrics=*/true, WarmupLevel::RAW, pool_, GetArrowPool(pool_))); // simulate read limits, only read 8 batches for (int32_t i = 0; i < 8; i++) { ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatchWithBitmap batch_with_bitmap, @@ -480,7 +474,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestReadWithoutInitializeReadRanges) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 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_))); + /*enable_io_metrics=*/false, WarmupLevel::RAW, pool_, GetArrowPool(pool_))); // simulate read limits, only read 8 batches ASSERT_NOK_WITH_MSG(reader->NextBatchWithBitmap(), "prefetch reader read ranges are not initialized"); @@ -498,7 +492,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestFailedIoMetrics) { /*prefetch_max_parallel_num=*/1, /*batch_size=*/10, /*prefetch_batch_count=*/2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/false, CacheConfig(), - /*enable_io_metrics=*/true, WarmupLevel::DECODED, pool_, GetArrowPool(pool_))); + /*enable_io_metrics=*/true, WarmupLevel::RAW, pool_, GetArrowPool(pool_))); ASSERT_NOK_WITH_MSG(reader->NextBatchWithBitmap(), "injected synchronous read failure"); std::shared_ptr metrics = reader->GetReaderMetrics(); @@ -745,7 +739,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, RefreshReadRanges) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 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_))); + /*enable_io_metrics=*/false, WarmupLevel::RAW, pool_, GetArrowPool(pool_))); auto prefetch_reader = dynamic_cast(reader.get()); ASSERT_OK(prefetch_reader->RefreshReadRanges()); std::vector> read_ranges_0 = {{0, 30}, {90, 101}}; @@ -776,7 +770,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, RefreshReadRangesDisablePrefetchByAdapti /*prefetch_batch_count=*/2, /*enable_adaptive_prefetch_strategy=*/true, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, WarmupLevel::DECODED, pool_, GetArrowPool(pool_))); + /*enable_io_metrics=*/false, WarmupLevel::RAW, pool_, GetArrowPool(pool_))); ASSERT_FALSE(reader->NeedPrefetch()); std::shared_ptr metrics = reader->GetReaderMetrics(); @@ -789,37 +783,35 @@ TEST_F(PrefetchFileBatchReaderImplTest, RefreshReadRangesDisablePrefetchByAdapti ASSERT_EQ(adaptive_disabled_count, 1); } -// DelegatingPrefetchReader is the layer that decides whether a read goes through the prefetch -// machinery or straight to the first inner reader, so its Warmup() has to follow the same decision. -// With prefetching on it must forward, otherwise a DECODED warmup silently degrades to a cold -// read. +// RAW warmup must reach the cache through the delegating layer without starting batch decoding. TEST_F(PrefetchFileBatchReaderImplTest, DelegatingPrefetchReaderWarmupForwardsWhenPrefetchEnabled) { auto data_array = PrepareArray(100); - int32_t batch_size = 10; - int32_t prefetch_max_parallel_num = 1; + const int32_t batch_size = 10; MockFormatReaderBuilder reader_builder(data_array, data_type_, batch_size); + reader_builder.SetPreBufferRangesPerSchema({{{0, 4}}}); ASSERT_OK_AND_ASSIGN( auto reader, PrefetchFileBatchReaderImpl::Create( - /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, - prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, - /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, WarmupLevel::DECODED, pool_, GetArrowPool(pool_))); + "", 0, &reader_builder, mock_fs_, /*prefetch_max_parallel_num=*/1, batch_size, + /*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::RAW, pool_, GetArrowPool(pool_))); auto* prefetch_reader = reader.get(); - ASSERT_TRUE(prefetch_reader->NeedPrefetch()); - DelegatingPrefetchReader delegating_reader(std::move(reader)); + ArrowSchema schema; + ASSERT_TRUE(arrow::ExportSchema(*arrow::schema(fields_), &schema).ok()); + ASSERT_OK(delegating_reader.SetReadSchema(&schema, nullptr, std::nullopt)); + ASSERT_TRUE(prefetch_reader->NeedPrefetch()); delegating_reader.Warmup(); - ASSERT_TRUE(DecodeThreadStarted(prefetch_reader)); - + ASSERT_FALSE(DecodeThreadStarted(prefetch_reader)); + ASSERT_OK_AND_ASSIGN(uint64_t io_count, delegating_reader.GetReaderMetrics()->GetCounter( + ReadAheadCacheMetrics::IO_COUNT)); + ASSERT_EQ(io_count, 1); ASSERT_OK_AND_ASSIGN(auto array, ReadResultCollector::CollectResult(&delegating_reader)); - auto expected_array = std::make_shared(data_array); - ASSERT_TRUE(expected_array->Equals(array)); + ASSERT_TRUE(std::make_shared(data_array)->Equals(array)); } -// Without prefetch, the synchronous wrapper prepares one batch and joins that task before the -// consumer reads the format reader again. The parallel Workloop must never start in this state. +// Format and adaptive fallback still warm raw bytes. The parallel Workloop never starts here. TEST_F(PrefetchFileBatchReaderImplTest, DelegatingPrefetchReaderWarmupWhenPrefetchBypassed) { for (bool format_prefetch : {false, true}) { auto data_array = PrepareArray(200); @@ -837,14 +829,22 @@ TEST_F(PrefetchFileBatchReaderImplTest, DelegatingPrefetchReaderWarmupWhenPrefet /*prefetch_batch_count=*/2, /*enable_adaptive_prefetch_strategy=*/true, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, WarmupLevel::DECODED, pool_, GetArrowPool(pool_))); + /*enable_io_metrics=*/false, WarmupLevel::RAW, pool_, GetArrowPool(pool_))); auto* prefetch_reader = reader.get(); ASSERT_FALSE(prefetch_reader->NeedPrefetch()); + auto* format_reader = dynamic_cast(prefetch_reader->GetFirstReader()); + ASSERT_NE(format_reader, nullptr); + format_reader->SetPreBufferRangesPerSchema({{{0, 4}}}); DelegatingPrefetchReader delegating_reader(std::move(reader)); + ArrowSchema schema; + ASSERT_TRUE(arrow::ExportSchema(*arrow::schema(fields_), &schema).ok()); + ASSERT_OK(delegating_reader.SetReadSchema(&schema, nullptr, std::nullopt)); delegating_reader.Warmup(); ASSERT_FALSE(DecodeThreadStarted(prefetch_reader)); - ASSERT_NE(prefetch_reader->synchronous_reader_->thread_, nullptr); + ASSERT_OK_AND_ASSIGN(uint64_t io_count, delegating_reader.GetReaderMetrics()->GetCounter( + ReadAheadCacheMetrics::IO_COUNT)); + ASSERT_EQ(io_count, 1); ASSERT_OK_AND_ASSIGN(auto array, ReadResultCollector::CollectResult(&delegating_reader)); auto expected_array = std::make_shared(data_array); @@ -868,7 +868,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, SetReadRanges) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 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_))); + /*enable_io_metrics=*/false, WarmupLevel::RAW, pool_, GetArrowPool(pool_))); auto prefetch_reader = dynamic_cast(reader.get()); ASSERT_FALSE(prefetch_reader->need_prefetch_); prefetch_reader->need_prefetch_ = true; @@ -912,7 +912,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, SetReadRangesReturnErrorWhenPushDownFail prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 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_))); + /*enable_io_metrics=*/false, WarmupLevel::RAW, pool_, GetArrowPool(pool_))); auto prefetch_reader = dynamic_cast(reader.get()); prefetch_reader->need_prefetch_ = true; @@ -938,8 +938,8 @@ TEST_F(PrefetchFileBatchReaderImplTest, WorkloopSetReadStatusWhenCacheInitFailed prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/true, - invalid_cache_config, /*enable_io_metrics=*/false, WarmupLevel::DECODED, - pool_, GetArrowPool(pool_))); + invalid_cache_config, /*enable_io_metrics=*/false, WarmupLevel::RAW, pool_, + GetArrowPool(pool_))); auto prefetch_reader = dynamic_cast(reader.get()); prefetch_reader->Workloop(); @@ -963,7 +963,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, DoReadBatchReturnOkWhenShutdown) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 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_))); + /*enable_io_metrics=*/false, WarmupLevel::RAW, pool_, GetArrowPool(pool_))); auto prefetch_reader = dynamic_cast(reader.get()); prefetch_reader->is_shutdown_ = true; @@ -982,7 +982,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, DoReadBatchReturnOkWhenNoCurrentReadRang prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 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_))); + /*enable_io_metrics=*/false, WarmupLevel::RAW, pool_, GetArrowPool(pool_))); auto prefetch_reader = dynamic_cast(reader.get()); prefetch_reader->read_ranges_in_group_ = {{}}; @@ -1001,7 +1001,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestReadWithLargeBatchSize) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, WarmupLevel::DECODED, pool_, GetArrowPool(pool_))); + /*enable_io_metrics=*/false, WarmupLevel::RAW, pool_, GetArrowPool(pool_))); ASSERT_NOK(reader->GetPreviousBatchFileRowId(0)); ASSERT_OK_AND_ASSIGN(auto array_and_row_ids, CollectResultAndRowIds(reader.get())); auto row_ids = array_and_row_ids.second; @@ -1022,7 +1022,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestPartialReaderSuccessRead) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, WarmupLevel::DECODED, pool_, GetArrowPool(pool_))); + /*enable_io_metrics=*/false, WarmupLevel::RAW, pool_, GetArrowPool(pool_))); auto prefetch_reader = dynamic_cast(reader.get()); for (int32_t i = 0; i < prefetch_max_parallel_num; i++) { dynamic_cast(prefetch_reader->readers_[i].get()) @@ -1068,7 +1068,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestAllReaderFailedWithIOError) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, WarmupLevel::DECODED, pool_, GetArrowPool(pool_))); + /*enable_io_metrics=*/false, WarmupLevel::RAW, pool_, GetArrowPool(pool_))); auto prefetch_reader = dynamic_cast(reader.get()); for (int32_t i = 0; i < prefetch_max_parallel_num; i++) { @@ -1107,7 +1107,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestPrefetchWithEmptyData) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, WarmupLevel::DECODED, pool_, GetArrowPool(pool_))); + /*enable_io_metrics=*/false, WarmupLevel::RAW, pool_, GetArrowPool(pool_))); ASSERT_NOK(reader->GetPreviousBatchFileRowId(0)); ASSERT_OK_AND_ASSIGN(auto array_and_row_ids, CollectResultAndRowIds(reader.get())); auto row_ids = array_and_row_ids.second; @@ -1127,7 +1127,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestCallNextBatchAfterReadingEof) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, WarmupLevel::DECODED, pool_, GetArrowPool(pool_))); + /*enable_io_metrics=*/false, WarmupLevel::RAW, pool_, GetArrowPool(pool_))); ASSERT_NOK(reader->GetPreviousBatchFileRowId(0)); ASSERT_OK_AND_ASSIGN(auto array_and_row_ids, CollectResultAndRowIds(reader.get())); auto row_ids = array_and_row_ids.second; @@ -1171,7 +1171,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestCreateReaderWithoutNextBatch) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, WarmupLevel::DECODED, pool_, GetArrowPool(pool_))); + /*enable_io_metrics=*/false, WarmupLevel::RAW, pool_, GetArrowPool(pool_))); } TEST_F(PrefetchFileBatchReaderImplTest, TestInvalidCase) { @@ -1186,7 +1186,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestInvalidCase) { /*prefetch_max_parallel_num=*/0, batch_size, 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, WarmupLevel::DECODED, pool_, GetArrowPool(pool_))); + /*enable_io_metrics=*/false, WarmupLevel::RAW, pool_, GetArrowPool(pool_))); } { ASSERT_NOK(PrefetchFileBatchReaderImpl::Create( @@ -1194,7 +1194,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestInvalidCase) { prefetch_max_parallel_num, /*batch_size=*/-1, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, WarmupLevel::DECODED, pool_, GetArrowPool(pool_))); + /*enable_io_metrics=*/false, WarmupLevel::RAW, pool_, GetArrowPool(pool_))); } { ASSERT_NOK_WITH_MSG( @@ -1203,7 +1203,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestInvalidCase) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, WarmupLevel::DECODED, pool_, GetArrowPool(pool_)), + /*enable_io_metrics=*/false, WarmupLevel::RAW, pool_, GetArrowPool(pool_)), "data file size should not be negative"); } { @@ -1213,7 +1213,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestInvalidCase) { /*enable_adaptive_prefetch_strategy=*/false, /*executor=*/nullptr, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), /*enable_io_metrics=*/false, - WarmupLevel::DECODED, pool_, GetArrowPool(pool_))); + WarmupLevel::RAW, pool_, GetArrowPool(pool_))); } { ASSERT_NOK(PrefetchFileBatchReaderImpl::Create( @@ -1221,7 +1221,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestInvalidCase) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, WarmupLevel::DECODED, pool_, GetArrowPool(pool_))); + /*enable_io_metrics=*/false, WarmupLevel::RAW, pool_, GetArrowPool(pool_))); } { ASSERT_NOK(PrefetchFileBatchReaderImpl::Create( @@ -1229,7 +1229,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestInvalidCase) { /*fs=*/nullptr, prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, WarmupLevel::DECODED, pool_, GetArrowPool(pool_))); + /*enable_io_metrics=*/false, WarmupLevel::RAW, pool_, GetArrowPool(pool_))); } { ASSERT_OK_AND_ASSIGN( @@ -1239,7 +1239,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestInvalidCase) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, WarmupLevel::DECODED, pool_, GetArrowPool(pool_))); + /*enable_io_metrics=*/false, WarmupLevel::RAW, pool_, GetArrowPool(pool_))); ASSERT_NOK_WITH_MSG(reader->SeekToRow(/*row_number=*/101), "not support seek to row for prefetch reader"); } @@ -1265,7 +1265,7 @@ TEST_P(PrefetchFileBatchReaderImplTest, TestPrefetchWithPredicatePushdownWithCom auto reader = PreparePrefetchReader(file_format, schema.get(), predicate, /*selection_bitmap=*/std::nullopt, /*batch_size=*/batch_size, /*prefetch_max_parallel_num=*/3, - read_ahead_cache_enabled, WarmupLevel::DECODED); + read_ahead_cache_enabled, WarmupLevel::RAW); ASSERT_OK_AND_ASSIGN(auto array_and_row_ids, CollectResultAndRowIds(reader.get())); arrow::ArrayVector expected_array_vector; @@ -1311,7 +1311,7 @@ TEST_P(PrefetchFileBatchReaderImplTest, auto reader = PreparePrefetchReader(file_format, schema.get(), predicate, /*selection_bitmap=*/std::nullopt, /*batch_size=*/batch_size, /*prefetch_max_parallel_num=*/3, - read_ahead_cache_enabled, WarmupLevel::DECODED); + read_ahead_cache_enabled, WarmupLevel::RAW); ASSERT_OK(reader->RefreshReadRanges()); ASSERT_NOK(reader->GetPreviousBatchFileRowId(0)); ASSERT_OK_AND_ASSIGN(auto array_and_row_ids, CollectResultAndRowIds(reader.get())); @@ -1343,7 +1343,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestPrefetchWithBitmap) { /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, WarmupLevel::DECODED, pool_, + /*enable_io_metrics=*/false, WarmupLevel::RAW, pool_, GetArrowPool(pool_))); ASSERT_OK_AND_ASSIGN(auto result_chunk_array, ReadResultCollector::CollectResult(std::move(reader))); @@ -1375,7 +1375,7 @@ TEST_P(PrefetchFileBatchReaderImplTest, TestRowMapping) { auto reader = PreparePrefetchReader(file_format, schema.get(), predicate, /*selection_bitmap=*/std::nullopt, /*batch_size=*/10, /*prefetch_max_parallel_num=*/3, - read_ahead_cache_enabled, WarmupLevel::DECODED); + read_ahead_cache_enabled, WarmupLevel::RAW); ASSERT_NOK(reader->GetPreviousBatchFileRowId(0)); ASSERT_OK_AND_ASSIGN(std::shared_ptr batch, paimon::test::ReadResultCollector::CollectResultOneBatch(reader.get())); @@ -1435,7 +1435,7 @@ TEST_P(PrefetchFileBatchReaderImplTest, FactoryWarmupIsIndependentOfPrefetch) { for (bool late_materializing : {false, true}) { for (int32_t mode = 0; mode < 3; ++mode) { std::vector baseline_row_ids; - for (WarmupLevel level : {WarmupLevel::NONE, WarmupLevel::RAW, WarmupLevel::DECODED}) { + for (WarmupLevel level : {WarmupLevel::NONE, WarmupLevel::RAW}) { SCOPED_TRACE(format + " mode=" + std::to_string(mode) + " level=" + std::to_string(static_cast(level)) + " late_materializing=" + std::to_string(late_materializing)); @@ -1504,7 +1504,7 @@ TEST_P(PrefetchFileBatchReaderImplTest, SynchronousMultiFileWarmupPreservesRows) PrepareTestData(format, data, /*stripe_row_count=*/30, /*row_index_stride=*/10); const std::string path = PathUtil::JoinPath(dir_->Str(), "file." + format); ASSERT_OK_AND_ASSIGN(auto file_status, local_fs_->GetFileStatus(path)); - for (WarmupLevel level : {WarmupLevel::NONE, WarmupLevel::RAW, WarmupLevel::DECODED}) { + for (WarmupLevel level : {WarmupLevel::NONE, WarmupLevel::RAW}) { DataFileReadOptions options; options.read_batch_size = 10; options.read_ahead_cache_enabled = cache_enabled; @@ -1563,8 +1563,7 @@ TEST_P(PrefetchFileBatchReaderImplTest, TestWarmupLevelNone) { } // WarmupLevel::RAW fetches the file's raw bytes into memory but does NOT start the decoder, so no -// batch is produced until the first real read. It overlaps the remote fetch while keeping memory -// lower than DECODED. +// batch is produced until the first real read. It can overlap the remote fetch without decoding. TEST_P(PrefetchFileBatchReaderImplTest, TestWarmupLevelRaw) { auto [file_format, read_ahead_cache_enabled] = GetParam(); auto data_array = PrepareArray(90); @@ -1640,37 +1639,6 @@ TEST_P(PrefetchFileBatchReaderImplTest, TestWarmupCacheRearmsForNewReadRangeGene ASSERT_TRUE(expected_array->Equals(array_and_row_ids.first)); } -// WarmupLevel::DECODED (the default) starts the background decode loop during Warmup(). Reading -// returns every row and reports produced batches, matching the behavior from before warmup levels -// existed. -TEST_P(PrefetchFileBatchReaderImplTest, TestWarmupLevelDecoded) { - auto [file_format, read_ahead_cache_enabled] = GetParam(); - auto data_array = PrepareArray(90); - PrepareTestData(file_format, data_array, /*stripe_row_count=*/30, /*row_index_stride=*/10); - auto schema = arrow::schema(fields_); - auto reader = PreparePrefetchReader(file_format, schema.get(), /*predicate=*/nullptr, - /*selection_bitmap=*/std::nullopt, /*batch_size=*/10, - /*prefetch_max_parallel_num=*/3, read_ahead_cache_enabled, - WarmupLevel::DECODED); - - reader->Warmup(); - - // The discriminating assertion: DECODED is the only level that starts the decode loop inside - // Warmup() itself, and it must have done so before the first read. Checking produced batches - // instead cannot tell DECODED from NONE, because the first NextBatch starts the same loop - // lazily and every level ends up with produced_batches > 0 once the read is over. - ASSERT_TRUE(DecodeThreadStarted(reader.get())); - - ASSERT_OK_AND_ASSIGN(auto array_and_row_ids, CollectResultAndRowIds(reader.get())); - auto expected_array = std::make_shared(data_array); - ASSERT_TRUE(expected_array->Equals(array_and_row_ids.first)); - - std::shared_ptr metrics = reader->GetReaderMetrics(); - ASSERT_OK_AND_ASSIGN(uint64_t produced_batches, - metrics->GetCounter(PrefetchMetrics::PRODUCED_BATCHES)); - ASSERT_GT(produced_batches, 0); -} - // The cache's stream used to be opened before the readers' streams were dispatched, which // serialized one remote round trip ahead of all of them. They are opened together now, and this // pins that down: the barrier needs every one of them in flight at once, so opening the cache's @@ -1697,7 +1665,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestCreateOpensAllStreamsConcurrently) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, WarmupLevel::DECODED, pool_, GetArrowPool(pool_))); + /*enable_io_metrics=*/false, WarmupLevel::RAW, pool_, GetArrowPool(pool_))); ASSERT_NE(reader, nullptr); ASSERT_TRUE(barrier->WaitForArrival()); } diff --git a/src/paimon/common/reader/warmup_file_batch_reader.cpp b/src/paimon/common/reader/warmup_file_batch_reader.cpp index f81646593..183edcd49 100644 --- a/src/paimon/common/reader/warmup_file_batch_reader.cpp +++ b/src/paimon/common/reader/warmup_file_batch_reader.cpp @@ -34,26 +34,13 @@ WarmupFileBatchReader::WarmupFileBatchReader(std::shared_ptrSetPreBufferRangeCallback(nullptr); read_ahead_->ReleaseBuffers(); } -void WarmupFileBatchReader::Join() const { - if (thread_) { - thread_->join(); - thread_.reset(); - } -} - -void WarmupFileBatchReader::ResetWarmup() { - Join(); - if (batch_ && batch_->ok()) { - ReaderUtils::ReleaseReadBatch(std::move(batch_.value()).value().first); - } - batch_.reset(); +void WarmupFileBatchReader::ResetReadState() { prepared_ = false; - started_ = false; batch_returned_ = false; read_status_ = Status::OK(); } @@ -63,15 +50,10 @@ void WarmupFileBatchReader::Prepare() { } void WarmupFileBatchReader::Warmup() { - if (!prepared_ || closed_ || started_ || level_ == WarmupLevel::NONE) { + if (!prepared_ || closed_ || level_ != WarmupLevel::RAW) { return; } - if (level_ == WarmupLevel::RAW) { - (void)read_ahead_->Warmup(reader_.get()); - } else { - started_ = true; - thread_ = std::make_unique([this]() { batch_.emplace(ReadOneBatch()); }); - } + (void)read_ahead_->Warmup(reader_.get()); } Result WarmupFileBatchReader::ReadOneBatch() { @@ -90,12 +72,8 @@ Result WarmupFileBatchReader::NextBatchWithBit return Status::Invalid("cannot read a closed warmup reader"); } PAIMON_RETURN_NOT_OK(read_status_); - // A normal read is synchronous unless Warmup() explicitly prepared the first batch. - started_ = true; - Join(); batch_returned_ = false; - Result result = batch_ ? std::move(batch_.value()) : ReadOneBatch(); - batch_.reset(); + Result result = ReadOneBatch(); if (result.ok() && !BatchReader::IsEofBatch(result.value())) { batch_returned_ = true; } else if (!result.ok()) { @@ -112,7 +90,7 @@ Result WarmupFileBatchReader::NextBatch() { Status WarmupFileBatchReader::SetReadSchema(::ArrowSchema* schema, const std::shared_ptr& predicate, const std::optional& bitmap) { - ResetWarmup(); + ResetReadState(); if (closed_) { return Status::Invalid("cannot configure a closed warmup reader"); } @@ -124,13 +102,8 @@ Status WarmupFileBatchReader::SetReadSchema(::ArrowSchema* schema, } std::shared_ptr WarmupFileBatchReader::GetReaderMetrics() const { - // The underlying reader may be decoding. Do not inspect its mutable metrics until it stops. - // Joining is reserved for operations needing the format reader, so a live snapshot contains - // only the thread-safe cache and I/O counters while the first batch is being prepared. auto metrics = std::make_shared(); - if (!thread_) { - metrics->Merge(reader_->GetReaderMetrics()); - } + metrics->Merge(reader_->GetReaderMetrics()); metrics->Merge(read_ahead_->GetMetrics()); metrics->SetGauge(PrefetchMetrics::ENABLED, 0.0); metrics->SetGauge(PrefetchMetrics::PARALLELISM, 1.0); @@ -138,7 +111,6 @@ std::shared_ptr WarmupFileBatchReader::GetReaderMetrics() const { } Result> WarmupFileBatchReader::GetFileSchema() const { - Join(); return reader_->GetFileSchema(); } @@ -150,12 +122,10 @@ Result WarmupFileBatchReader::GetPreviousBatchFileRowId(uint64_t batch } Result WarmupFileBatchReader::GetNumberOfRows() const { - Join(); return reader_->GetNumberOfRows(); } bool WarmupFileBatchReader::SupportPreciseBitmapSelection() const { - Join(); return reader_->SupportPreciseBitmapSelection(); } @@ -163,7 +133,7 @@ void WarmupFileBatchReader::Close() { if (closed_) { return; } - ResetWarmup(); + ResetReadState(); closed_ = true; reader_->SetPreBufferRangeCallback(nullptr); read_ahead_->ReleaseBuffers(); diff --git a/src/paimon/common/reader/warmup_file_batch_reader.h b/src/paimon/common/reader/warmup_file_batch_reader.h index 98ac5a6d8..9da5645f0 100644 --- a/src/paimon/common/reader/warmup_file_batch_reader.h +++ b/src/paimon/common/reader/warmup_file_batch_reader.h @@ -20,7 +20,6 @@ #include #include -#include #include "arrow/c/abi.h" #include "paimon/common/reader/file_read_ahead.h" @@ -32,7 +31,7 @@ class MemoryPool; namespace paimon { -/// A synchronous file reader that may prepare one batch before consumption starts. +/// A synchronous file reader with independent byte caching and raw warmup. class WarmupFileBatchReader : public FileBatchReader { public: WarmupFileBatchReader(std::shared_ptr reader, @@ -54,21 +53,17 @@ class WarmupFileBatchReader : public FileBatchReader { /// For a containing reader which configures the underlying format reader itself. void Prepare(); - void ResetWarmup(); + void ResetReadState(); private: - void Join() const; Result ReadOneBatch(); std::shared_ptr reader_; std::shared_ptr read_ahead_; const WarmupLevel level_; std::shared_ptr arrow_pool_; - mutable std::unique_ptr thread_; - std::optional> batch_; Status read_status_; bool prepared_ = false; - bool started_ = false; bool closed_ = false; bool batch_returned_ = false; }; diff --git a/src/paimon/common/reader/warmup_file_batch_reader_test.cpp b/src/paimon/common/reader/warmup_file_batch_reader_test.cpp index 4e86de443..80d8c4b9f 100644 --- a/src/paimon/common/reader/warmup_file_batch_reader_test.cpp +++ b/src/paimon/common/reader/warmup_file_batch_reader_test.cpp @@ -17,10 +17,7 @@ */ #include "paimon/common/reader/warmup_file_batch_reader.h" -#include -#include #include -#include #include #include #include @@ -54,13 +51,9 @@ class ObservedReader : public MockFileBatchReader { } Result NextBatchWithBitmap() override { - const uint32_t count = ++read_count; + ++read_count; last_read_thread = std::this_thread::get_id(); - auto result = Read(); - if (count == 1) { - first_read_done.set_value(); - } - return result; + return Read(); } Result>> PreBufferRange() override { @@ -78,8 +71,7 @@ class ObservedReader : public MockFileBatchReader { } } - std::atomic read_count{0}; - std::promise first_read_done; + uint32_t read_count = 0; std::thread::id last_read_thread; bool read_bytes = false; Status range_status; @@ -191,7 +183,7 @@ TEST_F(WarmupFileBatchReaderTest, RawWarmsBytesWithoutDecodingAndIsIdempotent) { reader->Warmup(); reader->Warmup(); ASSERT_EQ(fs->gated->AsyncReadCount(), 1); - ASSERT_EQ(observed_->read_count.load(), 0); + ASSERT_EQ(observed_->read_count, 0); fs->gated->ReleaseAll(); observed_->read_bytes = true; ASSERT_OK_AND_ASSIGN(auto batch, reader->NextBatchWithBitmap()); @@ -209,7 +201,7 @@ TEST_F(WarmupFileBatchReaderTest, NoneAndRawWithoutCacheDoNotStartWork) { ASSERT_OK_AND_ASSIGN(auto reader, Open(config.first, config.second)); ASSERT_OK(Configure(reader.get())); reader->Warmup(); - ASSERT_EQ(observed_->read_count.load(), 0); + ASSERT_EQ(observed_->read_count, 0); if (config.second) { ASSERT_EQ(Counter(reader.get(), ReadAheadCacheMetrics::IO_COUNT), 0); } @@ -221,48 +213,53 @@ TEST_F(WarmupFileBatchReaderTest, NoneAndRawWithoutCacheDoNotStartWork) { } } -TEST_F(WarmupFileBatchReaderTest, DecodedPreparesOneBatchThenReadsSynchronously) { - for (bool cache_enabled : {false, true}) { - ASSERT_OK_AND_ASSIGN(auto reader, Open(WarmupLevel::DECODED, cache_enabled)); - ASSERT_OK(Configure(reader.get())); - auto done = observed_->first_read_done.get_future(); - reader->Warmup(); - reader->Warmup(); - ASSERT_EQ(done.wait_for(std::chrono::seconds(5)), std::future_status::ready); - ASSERT_EQ(observed_->read_count.load(), 1); - ASSERT_NOK(reader->GetPreviousBatchFileRowId(0)); - ASSERT_OK_AND_ASSIGN(auto first, reader->NextBatchWithBitmap()); - ASSERT_NE(observed_->last_read_thread, std::this_thread::get_id()); - ASSERT_EQ(observed_->read_count.load(), 1); - ASSERT_OK_AND_ASSIGN(uint64_t row_id, reader->GetPreviousBatchFileRowId(1)); - ASSERT_EQ(row_id, 1); - ReaderUtils::ReleaseReadBatch(std::move(first.first)); - reader->Warmup(); - ASSERT_OK_AND_ASSIGN(auto second, reader->NextBatchWithBitmap()); - ASSERT_EQ(observed_->read_count.load(), 2); - ASSERT_EQ(observed_->last_read_thread, std::this_thread::get_id()); - ASSERT_OK_AND_ASSIGN(row_id, reader->GetPreviousBatchFileRowId(0)); - ASSERT_EQ(row_id, 2); - ReaderUtils::ReleaseReadBatch(std::move(second.first)); - ASSERT_OK_AND_ASSIGN(double enabled, - reader->GetReaderMetrics()->GetGauge(PrefetchMetrics::ENABLED)); - ASSERT_EQ(enabled, 0.0); +TEST_F(WarmupFileBatchReaderTest, WarmupNeverDecodesAndReadsStaySynchronous) { + for (WarmupLevel level : {WarmupLevel::NONE, WarmupLevel::RAW}) { + for (bool cache_enabled : {false, true}) { + ASSERT_OK_AND_ASSIGN(auto reader, Open(level, cache_enabled)); + ASSERT_OK(Configure(reader.get())); + reader->Warmup(); + reader->Warmup(); + ASSERT_EQ(observed_->read_count, 0); + ASSERT_NOK(reader->GetPreviousBatchFileRowId(0)); + ASSERT_OK_AND_ASSIGN(auto first, reader->NextBatchWithBitmap()); + ASSERT_EQ(observed_->last_read_thread, std::this_thread::get_id()); + ASSERT_EQ(observed_->read_count, 1); + ASSERT_OK_AND_ASSIGN(uint64_t row_id, reader->GetPreviousBatchFileRowId(1)); + ASSERT_EQ(row_id, 1); + ReaderUtils::ReleaseReadBatch(std::move(first.first)); + reader->Warmup(); + ASSERT_EQ(observed_->read_count, 1); + ASSERT_OK_AND_ASSIGN(auto second, reader->NextBatchWithBitmap()); + ASSERT_EQ(observed_->last_read_thread, std::this_thread::get_id()); + ASSERT_EQ(observed_->read_count, 2); + ASSERT_OK_AND_ASSIGN(row_id, reader->GetPreviousBatchFileRowId(0)); + ASSERT_EQ(row_id, 2); + ReaderUtils::ReleaseReadBatch(std::move(second.first)); + ASSERT_OK_AND_ASSIGN(double enabled, + reader->GetReaderMetrics()->GetGauge(PrefetchMetrics::ENABLED)); + ASSERT_EQ(enabled, 0.0); + } } } -TEST_F(WarmupFileBatchReaderTest, SchemaResetDiscardsFirstBatchAndRearmsWarmup) { - ASSERT_OK_AND_ASSIGN(auto reader, Open(WarmupLevel::DECODED, true)); +TEST_F(WarmupFileBatchReaderTest, SchemaResetRearmsRawWarmup) { + ASSERT_OK_AND_ASSIGN(auto reader, Open(WarmupLevel::RAW, true)); ASSERT_OK(Configure(reader.get())); reader->Warmup(); + ASSERT_EQ(Counter(reader.get(), ReadAheadCacheMetrics::IO_COUNT), 1); + ASSERT_EQ(observed_->read_count, 0); ASSERT_OK(Configure(reader.get())); - ASSERT_NOK(reader->GetPreviousBatchFileRowId(0)); + ASSERT_EQ(Counter(reader.get(), ReadAheadCacheMetrics::IO_COUNT), 0); reader->Warmup(); + ASSERT_EQ(Counter(reader.get(), ReadAheadCacheMetrics::IO_COUNT), 1); + ASSERT_EQ(observed_->read_count, 0); ASSERT_OK_AND_ASSIGN(auto result, ReadResultCollector::CollectResult(reader.get())); ASSERT_TRUE(result->Equals(std::make_shared(data_))); } -TEST_F(WarmupFileBatchReaderTest, WarmupFailureIsReturnedOnlyWhenConsumed) { - ASSERT_OK_AND_ASSIGN(auto reader, Open(WarmupLevel::DECODED, false)); +TEST_F(WarmupFileBatchReaderTest, ReadFailureIsTerminalUntilSchemaReset) { + ASSERT_OK_AND_ASSIGN(auto reader, Open(WarmupLevel::RAW, false)); ASSERT_OK(Configure(reader.get())); observed_->SetNextBatchStatus(Status::IOError("first batch failed")); reader->Warmup(); @@ -274,12 +271,12 @@ TEST_F(WarmupFileBatchReaderTest, WarmupFailureIsReturnedOnlyWhenConsumed) { ASSERT_TRUE(result->Equals(std::make_shared(data_))); } -TEST_F(WarmupFileBatchReaderTest, EmptyFileAndUnconsumedWarmupCloseSafely) { +TEST_F(WarmupFileBatchReaderTest, EmptyFileAndRawWarmupCloseSafely) { for (bool empty : {false, true}) { if (empty) { data_ = data_->Slice(0, 0); } - ASSERT_OK_AND_ASSIGN(auto reader, Open(WarmupLevel::DECODED, true)); + ASSERT_OK_AND_ASSIGN(auto reader, Open(WarmupLevel::RAW, true)); ASSERT_OK(Configure(reader.get())); reader->Warmup(); if (empty) { @@ -291,7 +288,7 @@ TEST_F(WarmupFileBatchReaderTest, EmptyFileAndUnconsumedWarmupCloseSafely) { reader->Warmup(); ASSERT_NOK(reader->NextBatchWithBitmap()); } - ASSERT_OK_AND_ASSIGN(auto reader, Open(WarmupLevel::DECODED, true)); + ASSERT_OK_AND_ASSIGN(auto reader, Open(WarmupLevel::RAW, true)); ASSERT_OK(Configure(reader.get())); reader->Warmup(); reader.reset(); @@ -303,7 +300,7 @@ TEST_F(WarmupFileBatchReaderTest, RawRangeAndAsyncFailuresAreDeferredUntilRead) ASSERT_OK(Configure(reader.get())); observed_->range_status = Status::IOError("range discovery failed"); reader->Warmup(); - ASSERT_EQ(observed_->read_count.load(), 0); + ASSERT_EQ(observed_->read_count, 0); ASSERT_NOK_WITH_MSG(reader->NextBatchWithBitmap(), "range discovery failed"); } auto fs = std::make_shared(); diff --git a/src/paimon/core/deletionvectors/apply_deletion_vector_batch_reader_test.cpp b/src/paimon/core/deletionvectors/apply_deletion_vector_batch_reader_test.cpp index 8384f485e..b6a34d44c 100644 --- a/src/paimon/core/deletionvectors/apply_deletion_vector_batch_reader_test.cpp +++ b/src/paimon/core/deletionvectors/apply_deletion_vector_batch_reader_test.cpp @@ -97,7 +97,7 @@ class ApplyDeletionVectorBatchReaderTest /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, WarmupLevel::DECODED, pool, arrow_pool)); + /*enable_io_metrics=*/false, WarmupLevel::RAW, pool, arrow_pool)); } else { file_batch_reader = std::make_unique(data, target_type_, batch_size); diff --git a/src/paimon/core/operation/read_context_test.cpp b/src/paimon/core/operation/read_context_test.cpp index 4fb5e00b4..b733a3a1f 100644 --- a/src/paimon/core/operation/read_context_test.cpp +++ b/src/paimon/core/operation/read_context_test.cpp @@ -154,7 +154,7 @@ TEST(ReadContextTest, TestSetContent) { builder.EnablePredicateFilter(true); builder.EnablePrefetch(true); builder.SetReadAheadCacheEnabled(false); - builder.SetWarmupLevel(WarmupLevel::DECODED); + builder.SetWarmupLevel(WarmupLevel::NONE); builder.SetPrefetchBatchCount(1200); builder.SetPrefetchMaxParallelNum(6); builder.EnableMultiThreadRowToBatch(true); @@ -180,7 +180,7 @@ TEST(ReadContextTest, TestSetContent) { ASSERT_TRUE(ctx->EnablePredicateFilter()); ASSERT_TRUE(ctx->EnablePrefetch()); ASSERT_FALSE(ctx->ReadAheadCacheEnabled()); - ASSERT_EQ(WarmupLevel::DECODED, ctx->GetWarmupLevel()); + ASSERT_EQ(WarmupLevel::NONE, ctx->GetWarmupLevel()); ASSERT_EQ(1200, ctx->GetPrefetchBatchCount()); ASSERT_EQ(6, ctx->GetPrefetchMaxParallelNum()); ASSERT_TRUE(ctx->EnableMultiThreadRowToBatch()); @@ -201,7 +201,7 @@ TEST(ReadContextTest, TestSetContent) { } TEST(ReadContextTest, TestSetWarmupLevel) { - for (WarmupLevel level : {WarmupLevel::NONE, WarmupLevel::RAW, WarmupLevel::DECODED}) { + for (WarmupLevel level : {WarmupLevel::NONE, WarmupLevel::RAW}) { ReadContextBuilder builder("table_root_path"); // The setter hands back the builder so it chains like every other setter on it. ASSERT_EQ(&builder, &builder.SetWarmupLevel(level)); diff --git a/src/paimon/core/table/system/system_table_test.cpp b/src/paimon/core/table/system/system_table_test.cpp index 7591e53bf..e86e7e98d 100644 --- a/src/paimon/core/table/system/system_table_test.cpp +++ b/src/paimon/core/table/system/system_table_test.cpp @@ -239,8 +239,8 @@ TEST(SystemTableTest, TestReadOptimizedSystemTablePathParsing) { // A system table builds a fresh ReadContext for the data table underneath it, starting from the // defaults, so a setting that is not copied across silently reverts: a caller that asked for NONE -// or DECODED would get RAW back. `$ro` has its own builder chain and `$audit_log` and `$binlog` -// share one, so all three are pinned here. +// would get RAW back. `$ro` has its own builder chain and `$audit_log` and `$binlog` share one, +// so all three are pinned here. TEST(SystemTableTest, TestNewReadPropagatesWarmupLevel) { std::map options = {{Options::FILE_SYSTEM, "local"}, {Options::FILE_FORMAT, "orc"}}; @@ -251,7 +251,7 @@ TEST(SystemTableTest, TestNewReadPropagatesWarmupLevel) { BinlogSystemTable binlog(/*fs=*/nullptr, "/tmp/table", table_schema, options); ReadOptimizedSystemTable read_optimized("/tmp/table", table_schema, options); - for (WarmupLevel level : {WarmupLevel::NONE, WarmupLevel::RAW, WarmupLevel::DECODED}) { + for (WarmupLevel level : {WarmupLevel::NONE, WarmupLevel::RAW}) { ReadContextBuilder builder("/tmp/table"); builder.SetOptions(options).SetWarmupLevel(level); ASSERT_OK_AND_ASSIGN(std::unique_ptr caller_unique_context, builder.Finish()); From a35907d18e8fa686f59800506d7959e7121e41f4 Mon Sep 17 00:00:00 2001 From: "yonghao.fyh" Date: Thu, 5 Nov 2026 15:13:13 +0800 Subject: [PATCH 3/6] fix(read)!: bound warmup caching and share its context Detect synchronous byte-range support from reader capabilities and build once. Share the internal cache context with LM and ORC instead of range callbacks. Bound synchronous buffer residency, validate ranges before splitting, and discover ORC stripe ranges on demand while preserving Parquet buffer handoff. Cover lifecycle, adaptive fallback, payload warmup and high file positions. BREAKING CHANGE: remove PrefetchFileBatchReader::SetPreBufferRangeCallback and PreBufferRangeCallback. Plugins must remove overrides and rebuild. --- docs/source/user_guide/prefetch.rst | 24 +- include/paimon/read_context.h | 3 +- .../reader/prefetch_file_batch_reader.h | 15 - src/paimon/common/io/cache_input_stream.h | 43 ++- .../common/io/cache_input_stream_test.cpp | 35 ++ .../reader/data_file_reader_factory.cpp | 31 +- .../reader/data_file_reader_factory_test.cpp | 158 ++++++++- src/paimon/common/reader/file_read_ahead.cpp | 133 +++++-- src/paimon/common/reader/file_read_ahead.h | 32 +- .../late_materializing_file_batch_reader.cpp | 32 +- .../late_materializing_file_batch_reader.h | 34 +- ...e_materializing_file_batch_reader_test.cpp | 332 ++++++++++++++---- .../late_materializing_reader_builder.h | 9 +- .../prefetch_file_batch_reader_impl.cpp | 13 +- .../reader/prefetch_file_batch_reader_impl.h | 3 - .../reader/warmup_file_batch_reader.cpp | 2 - .../reader/warmup_file_batch_reader_test.cpp | 112 +++++- src/paimon/common/utils/read_ahead_cache.cpp | 93 ++++- src/paimon/common/utils/read_ahead_cache.h | 18 +- .../common/utils/read_ahead_cache_test.cpp | 141 ++++++++ .../format/orc/orc_file_batch_reader.cpp | 31 +- src/paimon/format/orc/orc_file_batch_reader.h | 14 +- .../format/orc/orc_file_batch_reader_test.cpp | 204 +++++++++++ src/paimon/format/orc/orc_reader_builder.h | 9 +- src/paimon/format/orc/orc_reader_wrapper.h | 42 ++- .../parquet_file_batch_reader_test.cpp | 40 +++ .../format/parquet/parquet_reader_builder.h | 21 +- test/inte/pk_compaction_inte_test.cpp | 4 - 28 files changed, 1360 insertions(+), 268 deletions(-) diff --git a/docs/source/user_guide/prefetch.rst b/docs/source/user_guide/prefetch.rst index d9595b09d..1eb96b189 100644 --- a/docs/source/user_guide/prefetch.rst +++ b/docs/source/user_guide/prefetch.rst @@ -92,8 +92,11 @@ context and always uses the default level. read-ahead cache without decoding them. It has no effect when the read-ahead cache is disabled with ``ReadContextBuilder::SetReadAheadCacheEnabled()``. -Warmup and the read-ahead cache work independently of batch prefetch for -Parquet and ORC. With ``EnablePrefetch(false)``, ``RAW`` still warms the cache +Warmup and the read-ahead cache work independently of batch prefetch for readers +implementing the ``PrefetchFileBatchReader`` byte-range interface, including +Parquet and ORC. Support is determined by the reader's capabilities rather than +the format name; other readers keep their plain read path. With +``EnablePrefetch(false)``, ``RAW`` still warms the cache and batches are read synchronously. The same behavior applies when the format or the adaptive strategy decides against prefetching. Warmup only fetches raw bytes; it never decodes batches. When batch prefetch is active, the first @@ -119,3 +122,20 @@ on top of the file being read. It may also fetch data from a file that the reader never consumes: a read that stops early, for example because of a LIMIT, may already have warmed the next file, and closing it waits for that file's in-flight fetches. + +For synchronous reads, including adaptive fallback, the pre-buffer limit also +bounds retained range buffers. Completed buffers outside the current window can +be evicted; later reads can fetch them again. In-flight I/O and buffers retained +by callbacks count against that budget, so a full budget can temporarily cause a +read to bypass range caching. The block cache has its separate limit. Parallel +prefetch retains its existing range-buffer retention policy. + +ORC discovers byte ranges one stripe at a time and reports later stripes as +reading advances. Preparing the first batch does not require warmup to parse +every stripe footer. + +Late-materialization payloads and later ORC stripes prepare their ranges through +the same internal cache context. Format readers expose their initial ranges +through ``PreBufferRange()``; the public prefetch reader has no range callback +setter. Plugins overriding the former callback setter must remove that override +and rebuild against the updated interface. diff --git a/include/paimon/read_context.h b/include/paimon/read_context.h index 20dab5902..cb06e05a7 100644 --- a/include/paimon/read_context.h +++ b/include/paimon/read_context.h @@ -414,7 +414,8 @@ class PAIMON_EXPORT ReadContextBuilder { /// /// A read-ahead cache is used to prebuffer data ranges before they are needed, /// which can improve read performance by reducing redundant I/O operations. - /// For Parquet and ORC it also works when batch prefetch is disabled. + /// It also works without batch prefetch for readers exposing byte ranges through + /// PrefetchFileBatchReader, including Parquet and ORC. /// @param enabled Whether to enable the read-ahead cache (default: true) /// @return Reference to this builder for method chaining. ReadContextBuilder& SetReadAheadCacheEnabled(bool enabled); diff --git a/include/paimon/reader/prefetch_file_batch_reader.h b/include/paimon/reader/prefetch_file_batch_reader.h index ed0ce55d5..eac8d046e 100644 --- a/include/paimon/reader/prefetch_file_batch_reader.h +++ b/include/paimon/reader/prefetch_file_batch_reader.h @@ -19,7 +19,6 @@ #pragma once -#include #include #include #include @@ -137,20 +136,6 @@ 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/src/paimon/common/io/cache_input_stream.h b/src/paimon/common/io/cache_input_stream.h index 015b082c3..e828e9947 100644 --- a/src/paimon/common/io/cache_input_stream.h +++ b/src/paimon/common/io/cache_input_stream.h @@ -18,6 +18,7 @@ #pragma once +#include #include #include @@ -27,11 +28,49 @@ namespace paimon { +class FileReadAhead; + class CacheInputStream : public InputStream { public: + using CacheFactory = std::function>()>; + CacheInputStream(std::unique_ptr input_stream, const std::shared_ptr& cache) - : cache_(cache), input_stream_(std::move(input_stream)) {} + : CacheInputStream(std::move(input_stream), cache, CacheFactory(), nullptr) {} + + CacheInputStream(std::unique_ptr input_stream, + const std::shared_ptr& cache, CacheFactory cache_factory) + : CacheInputStream(std::move(input_stream), cache, std::move(cache_factory), nullptr) {} + + CacheInputStream(std::unique_ptr input_stream, + const std::shared_ptr& cache, CacheFactory cache_factory, + const std::shared_ptr& read_ahead) + : cache_(cache), + input_stream_(std::move(input_stream)), + cache_factory_(std::move(cache_factory)), + read_ahead_(read_ahead) {} + + const std::shared_ptr& GetReadAhead() const { + return read_ahead_; + } + + // Configure before the reader starts consuming data. Capable builders may enable the cache + // early to coordinate their own buffering; otherwise the factory enables it after Build(). + Result EnableCache() { + if (cache_) { + return true; + } + if (!cache_factory_) { + return false; + } + PAIMON_ASSIGN_OR_RAISE(cache_, cache_factory_()); + cache_factory_ = nullptr; + return cache_ != nullptr; + } + + void ClearCacheFactory() { + cache_factory_ = nullptr; + } Status Seek(int64_t offset, SeekOrigin origin) override { return input_stream_->Seek(offset, origin); @@ -96,6 +135,8 @@ class CacheInputStream : public InputStream { private: std::shared_ptr cache_; std::unique_ptr input_stream_; + CacheFactory cache_factory_; + std::shared_ptr read_ahead_; }; } // namespace paimon diff --git a/src/paimon/common/io/cache_input_stream_test.cpp b/src/paimon/common/io/cache_input_stream_test.cpp index 4c05c57c3..52a0883b9 100644 --- a/src/paimon/common/io/cache_input_stream_test.cpp +++ b/src/paimon/common/io/cache_input_stream_test.cpp @@ -26,6 +26,7 @@ #include "gtest/gtest.h" #include "paimon/common/factories/io_hook.h" +#include "paimon/common/metrics/metrics_impl.h" #include "paimon/common/utils/read_ahead_cache.h" #include "paimon/common/utils/scope_guard.h" #include "paimon/fs/file_system.h" @@ -82,6 +83,40 @@ class CacheInputStreamTest : public ::testing::Test { std::string file_path_; }; +TEST_F(CacheInputStreamTest, EnableCacheLazilyOnce) { + uint32_t factory_calls = 0; + auto cache = CreateCache({{2, 5}}); + CacheInputStream stream(OpenFile(), nullptr, [&]() -> Result> { + ++factory_calls; + return cache; + }); + std::string buffer(5, '\0'); + ASSERT_OK(stream.Read(buffer.data(), 5, 2)); + ASSERT_EQ(factory_calls, 0); + ASSERT_EQ(buffer, "cdefg"); + ASSERT_OK_AND_ASSIGN(bool enabled, stream.EnableCache()); + ASSERT_TRUE(enabled); + ASSERT_OK_AND_ASSIGN(bool enabled_again, stream.EnableCache()); + ASSERT_TRUE(enabled_again); + ASSERT_EQ(factory_calls, 1); + ASSERT_OK(stream.Read(buffer.data(), 5, 2)); + ASSERT_EQ(buffer, "cdefg"); + std::shared_ptr metrics = std::make_shared(); + cache->CollectMetrics(&metrics); + ASSERT_OK_AND_ASSIGN(uint64_t io_count, metrics->GetCounter(ReadAheadCacheMetrics::IO_COUNT)); + ASSERT_EQ(io_count, 1); +} + +TEST_F(CacheInputStreamTest, DisableCacheFactoryAndPropagateFailure) { + CacheInputStream stream(OpenFile(), nullptr, []() -> Result> { + return Status::IOError("cache initialization failed"); + }); + ASSERT_NOK_WITH_MSG(stream.EnableCache(), "cache initialization failed"); + stream.ClearCacheFactory(); + ASSERT_OK_AND_ASSIGN(bool enabled, stream.EnableCache()); + ASSERT_FALSE(enabled); +} + // Test proxy methods: Seek, GetPos, Read (sequential), Close, GetUri, Length TEST_F(CacheInputStreamTest, TestProxyMethods) { auto underlying = OpenFile(); diff --git a/src/paimon/common/reader/data_file_reader_factory.cpp b/src/paimon/common/reader/data_file_reader_factory.cpp index c53ab943a..2bba5fff8 100644 --- a/src/paimon/common/reader/data_file_reader_factory.cpp +++ b/src/paimon/common/reader/data_file_reader_factory.cpp @@ -24,7 +24,6 @@ #include "paimon/common/reader/delegating_prefetch_reader.h" #include "paimon/common/reader/file_read_ahead.h" #include "paimon/common/reader/prefetch_file_batch_reader_impl.h" -#include "paimon/common/reader/warmup_file_batch_reader.h" #include "paimon/format/file_format.h" #include "paimon/format/file_format_factory.h" #include "paimon/format/read_hints.h" @@ -48,11 +47,6 @@ bool FormatSupportsPrefetch(const std::string& format_identifier) { format_identifier != "vortex"; } -bool FormatSupportsReadAhead(const std::string& format_identifier) { - // Other format plugins may expose only FileBatchReader. Keep their plain read path intact. - return format_identifier == "parquet" || format_identifier == "orc"; -} - } // namespace Result> DataFileReaderFactory::CreateReaderBuilder( @@ -76,10 +70,9 @@ Result> DataFileReaderFactory::CreateReaderBuilde // off when the shared read-ahead cache is reading ahead. ReadHints read_hints; read_hints.prefetch_enabled = read_options.prefetch_enabled; - read_hints.read_ahead_cache_enabled = - read_options.read_ahead_cache_enabled && - ((read_options.prefetch_enabled && FormatSupportsPrefetch(format_identifier)) || - FormatSupportsReadAhead(format_identifier)); + read_hints.read_ahead_cache_enabled = read_options.read_ahead_cache_enabled && + read_options.prefetch_enabled && + FormatSupportsPrefetch(format_identifier); reader_builder->WithReadHints(read_hints); return reader_builder; } @@ -102,20 +95,10 @@ Result> DataFileReaderFactory::Open( read_options.warmup_level, pool, arrow_pool)); return std::make_unique(std::move(prefetch_reader)); } - if (FormatSupportsReadAhead(format_identifier)) { - PAIMON_ASSIGN_OR_RAISE( - FileReadAhead::OpenedReaders opened, - FileReadAhead::Open(file_path, file_size, reader_builder, file_system, - /*reader_count=*/1, /*executor=*/nullptr, - read_options.read_ahead_cache_enabled, read_options.cache_config, - read_options.prefetch_io_metrics_enabled, pool)); - return std::make_unique(std::move(opened.readers[0]), - std::move(opened.read_ahead), - read_options.warmup_level, arrow_pool); - } - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr input_stream, - file_system->Open(FileStatus(file_path, file_size))); - return reader_builder->Build(input_stream); + return FileReadAhead::OpenSynchronous( + file_path, file_size, reader_builder, file_system, read_options.read_ahead_cache_enabled, + read_options.cache_config, read_options.prefetch_io_metrics_enabled, + read_options.warmup_level, pool, arrow_pool); } } // namespace paimon diff --git a/src/paimon/common/reader/data_file_reader_factory_test.cpp b/src/paimon/common/reader/data_file_reader_factory_test.cpp index 65221bd42..dd5c0ffb5 100644 --- a/src/paimon/common/reader/data_file_reader_factory_test.cpp +++ b/src/paimon/common/reader/data_file_reader_factory_test.cpp @@ -24,8 +24,11 @@ #include "arrow/array/array_nested.h" #include "arrow/ipc/json_simple.h" #include "gtest/gtest.h" +#include "paimon/common/io/cache_input_stream.h" #include "paimon/common/reader/delegating_prefetch_reader.h" +#include "paimon/common/reader/warmup_file_batch_reader.h" #include "paimon/common/utils/arrow/mem_utils.h" +#include "paimon/common/utils/checked_cast.h" #include "paimon/executor.h" #include "paimon/memory/memory_pool.h" #include "paimon/reader/file_batch_reader.h" @@ -36,21 +39,166 @@ namespace paimon::test { -TEST(DataFileReaderFactoryTest, OtherFormatsKeepTheirPlainReaderWithoutPrefetch) { +// Deliberately expose only FileBatchReader, even though the test data source can report ranges. +class PlainFileBatchReader : public FileBatchReader { + public: + explicit PlainFileBatchReader(std::unique_ptr reader) + : reader_(std::move(reader)) {} + + Result NextBatch() override { + return reader_->NextBatch(); + } + std::shared_ptr GetReaderMetrics() const override { + return reader_->GetReaderMetrics(); + } + Result> GetFileSchema() const override { + return reader_->GetFileSchema(); + } + Status SetReadSchema(::ArrowSchema* schema, const std::shared_ptr& predicate, + const std::optional& bitmap) override { + return reader_->SetReadSchema(schema, predicate, bitmap); + } + Result GetPreviousBatchFileRowId(uint64_t row) const override { + return reader_->GetPreviousBatchFileRowId(row); + } + Result GetNumberOfRows() const override { + return reader_->GetNumberOfRows(); + } + bool SupportPreciseBitmapSelection() const override { + return reader_->SupportPreciseBitmapSelection(); + } + void Close() override { + reader_->Close(); + } + + private: + std::unique_ptr reader_; +}; + +class CapabilityReaderBuilder : public MockFormatReaderBuilder { + public: + CapabilityReaderBuilder(const std::shared_ptr& data, + const std::shared_ptr& type, bool capable) + : MockFormatReaderBuilder(data, type, /*read_batch_size=*/2), capable_(capable) { + SetPreBufferRangesPerSchema({{{0, 4}}}); + } + + Result> Build( + const std::shared_ptr& input) const override { + ++build_count; + stream = std::dynamic_pointer_cast(input); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader, + MockFormatReaderBuilder::Build(input)); + checked_cast(reader.get())->EnableRandomizeBatchSize(false); + if (!capable_) { + reader = std::make_unique(std::move(reader)); + } + original_reader = reader.get(); + return reader; + } + + mutable uint32_t build_count = 0; + mutable FileBatchReader* original_reader = nullptr; + mutable std::shared_ptr stream; + + private: + bool capable_; +}; + +class CountingOpenFileSystem : public MockFileSystem { + public: + using MockFileSystem::Open; + Result> Open(const std::string& path) const override { + if (++open_count == 2 && fail_cache_open) { + return Status::IOError("cache stream open failed"); + } + return MockFileSystem::Open(path); + } + mutable uint32_t open_count = 0; + bool fail_cache_open = false; +}; + +TEST(DataFileReaderFactoryTest, PlainReaderCapabilityOverridesFormatName) { auto type = arrow::struct_({arrow::field("v", arrow::int32())}); auto data = arrow::ipc::internal::json::ArrayFromJSON(type, "[[1],[2]]").ValueOrDie(); - MockFormatReaderBuilder builder(data, type, /*read_batch_size=*/2); - auto fs = std::make_shared(); auto pool = GetDefaultPool(); DataFileReadOptions options; - for (const char* format : {"custom", "avro", "blob", "lance", "mosaic", "vortex"}) { + for (const char* format : + {"custom", "parquet", "orc", "avro", "blob", "lance", "mosaic", "vortex"}) { + SCOPED_TRACE(format); + CapabilityReaderBuilder builder(data, type, /*capable=*/false); + auto fs = std::make_shared(); + fs->fail_cache_open = true; ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, DataFileReaderFactory::Open(format, "dummy", 0, &builder, options, fs, nullptr, pool, GetArrowPool(pool))); - ASSERT_NE(dynamic_cast(reader.get()), nullptr); + ASSERT_EQ(reader.get(), builder.original_reader); + ASSERT_EQ(builder.build_count, 1); + ASSERT_EQ(fs->open_count, 1); + ASSERT_NE(builder.stream, nullptr); + ASSERT_OK_AND_ASSIGN(bool cache_active, builder.stream->EnableCache()); + ASSERT_FALSE(cache_active); + reader->Warmup(); + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, reader->NextBatch()); + ASSERT_EQ(batch.first->length, 2); + batch.first->release(batch.first.get()); + batch.second->release(batch.second.get()); + ASSERT_EQ(fs->open_count, 1); } } +TEST(DataFileReaderFactoryTest, CapableReaderSupportsReadAheadRegardlessOfFormatName) { + auto type = arrow::struct_({arrow::field("v", arrow::int32())}); + auto data = arrow::ipc::internal::json::ArrayFromJSON(type, "[[1],[2]]").ValueOrDie(); + auto pool = GetDefaultPool(); + for (const char* format : {"custom", "avro", "parquet", "orc"}) { + for (bool cache_enabled : {false, true}) { + SCOPED_TRACE(format); + SCOPED_TRACE(cache_enabled); + CapabilityReaderBuilder builder(data, type, /*capable=*/true); + auto fs = std::make_shared(); + DataFileReadOptions options; + options.read_ahead_cache_enabled = cache_enabled; + options.prefetch_max_parallel_num = 0; + options.prefetch_batch_count = 0; + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + DataFileReaderFactory::Open(format, "dummy", 0, &builder, options, fs, nullptr, + pool, GetArrowPool(pool))); + ASSERT_NE(dynamic_cast(reader.get()), nullptr); + ASSERT_EQ(builder.build_count, 1); + ASSERT_EQ(fs->open_count, cache_enabled ? 2 : 1); + ASSERT_OK_AND_ASSIGN(std::unique_ptr schema, reader->GetFileSchema()); + ASSERT_OK(reader->SetReadSchema(schema.get(), nullptr, std::nullopt)); + reader->Warmup(); + if (cache_enabled) { + ASSERT_OK_AND_ASSIGN(uint64_t io_count, reader->GetReaderMetrics()->GetCounter( + ReadAheadCacheMetrics::IO_COUNT)); + ASSERT_EQ(io_count, 1); + } + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, reader->NextBatch()); + ASSERT_EQ(batch.first->length, 2); + batch.first->release(batch.first.get()); + batch.second->release(batch.second.get()); + reader->Close(); + } + } +} + +TEST(DataFileReaderFactoryTest, CapableReaderPropagatesCacheOpenFailureWithoutRebuilding) { + auto type = arrow::struct_({arrow::field("v", arrow::int32())}); + auto data = arrow::ipc::internal::json::ArrayFromJSON(type, "[[1],[2]]").ValueOrDie(); + CapabilityReaderBuilder builder(data, type, /*capable=*/true); + auto fs = std::make_shared(); + fs->fail_cache_open = true; + auto pool = GetDefaultPool(); + ASSERT_NOK_WITH_MSG(DataFileReaderFactory::Open("custom", "dummy", 0, &builder, {}, fs, nullptr, + pool, GetArrowPool(pool)), + "cache stream open failed"); + ASSERT_EQ(builder.build_count, 1); + ASSERT_EQ(fs->open_count, 2); +} + /// `Open()` is the one place a production read builds the prefetching reader, and a level that is /// not forwarded here fails nothing: the reader warms up however its own default says, and the only /// symptom is latency and memory a query never asked for. diff --git a/src/paimon/common/reader/file_read_ahead.cpp b/src/paimon/common/reader/file_read_ahead.cpp index f48aafb72..5322363b3 100644 --- a/src/paimon/common/reader/file_read_ahead.cpp +++ b/src/paimon/common/reader/file_read_ahead.cpp @@ -24,9 +24,12 @@ #include #include +#include "fmt/format.h" #include "paimon/common/executor/future.h" #include "paimon/common/io/cache_input_stream.h" #include "paimon/common/metrics/metrics_impl.h" +#include "paimon/common/reader/warmup_file_batch_reader.h" +#include "paimon/common/utils/checked_cast.h" #include "paimon/format/reader_builder.h" namespace paimon { @@ -163,6 +166,11 @@ class MetricsInputStream : public InputStream { } // namespace +std::shared_ptr FileReadAhead::Create(const std::shared_ptr& cache, + const std::string& path) { + return std::shared_ptr(new FileReadAhead(cache, path)); +} + Result FileReadAhead::Open( const std::string& path, int64_t file_size, const ReaderBuilder* builder, const std::shared_ptr& fs, uint32_t reader_count, @@ -221,6 +229,8 @@ Result FileReadAhead::Open( cache = std::make_shared(input_stream, cache_config, static_cast(file_size), pool); } + auto state = Create(cache, path); + state->io_metrics_ = io_metrics; // Wave two builds the readers, which reads each file's footer and so still needs the cache // above to exist. The builds are concurrent, like the opens were. std::vector>>> futures; @@ -228,15 +238,15 @@ Result FileReadAhead::Open( for (uint32_t i = 0; i < reader_count; i++) { futures.push_back(RunOrSubmit( executor.get(), - [&builder, &cache, &streams, i, next_stream, - io_metrics]() -> Result> { + [&builder, &cache, &streams, i, next_stream, io_metrics, + state]() -> Result> { std::unique_ptr input_stream = std::move(streams[next_stream + i]); if (io_metrics) { input_stream = std::make_unique(std::move(input_stream), io_metrics); } - auto cache_input_stream = - std::make_shared(std::move(input_stream), cache); + auto cache_input_stream = std::make_shared( + std::move(input_stream), cache, CacheInputStream::CacheFactory(), state); return builder->Build(cache_input_stream); })); } @@ -254,26 +264,92 @@ Result FileReadAhead::Open( } readers.emplace_back(prefetch_file_batch_reader); } - auto state = std::shared_ptr(new FileReadAhead(cache, io_metrics)); - for (const auto& reader : readers) { - std::weak_ptr weak_state = state; - reader->SetPreBufferRangeCallback( - [weak_state](std::vector>&& ranges) { - if (auto shared_state = weak_state.lock()) { - shared_state->RegisterLateRanges(std::move(ranges)); - } - }); - } return OpenedReaders{std::move(readers), std::move(state)}; } +Result> FileReadAhead::OpenSynchronous( + const std::string& path, int64_t file_size, const ReaderBuilder* builder, + const std::shared_ptr& fs, bool cache_enabled, const CacheConfig& cache_config, + bool enable_io_metrics, WarmupLevel level, const std::shared_ptr& pool, + const std::shared_ptr& arrow_pool) { + if (file_size < 0 || !builder || !fs) { + return Status::Invalid("invalid synchronous file read-ahead open arguments"); + } + std::shared_ptr io_metrics; + if (enable_io_metrics) { + io_metrics = std::make_shared(); + } + auto state = Create(nullptr, path); + state->io_metrics_ = io_metrics; + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr input_stream, + fs->Open(FileStatus(path, file_size))); + if (io_metrics) { + input_stream = std::make_unique(std::move(input_stream), io_metrics); + } + CacheInputStream::CacheFactory cache_factory; + if (cache_enabled) { + cache_factory = [state, fs, path, file_size, cache_config, + pool]() -> Result> { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr cache_stream, + fs->Open(FileStatus(path, file_size))); + if (state->io_metrics_) { + cache_stream = + std::make_shared(cache_stream, state->io_metrics_); + } + state->cache_ = std::make_shared( + cache_stream, cache_config, static_cast(file_size), pool); + return state->cache_; + }; + } + auto stream = std::make_shared(std::move(input_stream), nullptr, + std::move(cache_factory), state); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader, builder->Build(stream)); + // The byte-range interface, rather than the format name, determines cache/warmup support. + if (!dynamic_cast(reader.get())) { + stream->ClearCacheFactory(); + return reader; + } + PAIMON_RETURN_NOT_OK(stream->EnableCache().status()); + state->SetSynchronous(true); + std::shared_ptr capable_reader = + checked_pointer_cast(std::move(reader)); + return std::make_unique(std::move(capable_reader), std::move(state), + level, arrow_pool); +} + Status FileReadAhead::Warmup(PrefetchFileBatchReader* reader) { std::lock_guard lock(mutex_); if (!cache_ || warmed_) { return status_; } warmed_ = true; - auto ranges = reader->PreBufferRange(); + return PrepareRangesLocked(reader->PreBufferRange(), true); +} + +Status FileReadAhead::PrepareRanges(PrefetchFileBatchReader* reader) { + { + std::lock_guard lock(mutex_); + if (!cache_ || !warmed_) { + return Status::OK(); + } + if (!status_.ok()) { + return status_; + } + } + // Range discovery may read remote metadata. Keep it outside the context lock so parallel + // format readers can discover their stripes concurrently. Reset/Close stop reader tasks first. + Result ranges = reader->PreBufferRange(); + std::lock_guard lock(mutex_); + if (!cache_ || !warmed_) { + return Status::OK(); + } + if (!status_.ok()) { + return status_; + } + return PrepareRangesLocked(std::move(ranges), false); +} + +Status FileReadAhead::PrepareRangesLocked(Result&& ranges, bool initial) { if (!ranges.ok()) { status_ = ranges.status(); ++error_count_; @@ -285,29 +361,22 @@ Status FileReadAhead::Warmup(PrefetchFileBatchReader* reader) { } auto added = cache_->AddRanges(std::move(byte_ranges)); if (!added.ok()) { - status_ = added.status(); + status_ = Status(added.status().code(), fmt::format("read-ahead ranges for file '{}': {}", + path_, added.status().ToString())); ++error_count_; return status_; } - cache_->Warmup(); + if (initial) { + cache_->Warmup(); + } else if (added.value()) { + cache_->Warmup(added.value().value()); + } return Status::OK(); } -void FileReadAhead::RegisterLateRanges(std::vector>&& ranges) { - std::lock_guard lock(mutex_); - if (!cache_ || !warmed_ || ranges.empty() || !status_.ok()) { - return; - } - std::vector byte_ranges; - for (const auto& range : ranges) { - byte_ranges.emplace_back(range.first, range.second); - } - auto added = cache_->AddRanges(std::move(byte_ranges)); - if (!added.ok()) { - status_ = added.status(); - ++error_count_; - } else if (added.value()) { - cache_->Warmup(added.value().value()); +void FileReadAhead::SetSynchronous(bool synchronous) { + if (cache_) { + cache_->SetBoundedRetention(synchronous); } } diff --git a/src/paimon/common/reader/file_read_ahead.h b/src/paimon/common/reader/file_read_ahead.h index 0979b994c..4a4dfa4b5 100644 --- a/src/paimon/common/reader/file_read_ahead.h +++ b/src/paimon/common/reader/file_read_ahead.h @@ -27,6 +27,11 @@ #include "paimon/common/utils/read_ahead_cache.h" #include "paimon/reader/prefetch_file_batch_reader.h" +#include "paimon/visibility.h" + +namespace arrow { +class MemoryPool; +} // namespace arrow namespace paimon { @@ -34,15 +39,19 @@ class Executor; class ReaderBuilder; struct PrefetchIoMetricsState; -/// Per-file byte caching and I/O instrumentation, shared by synchronous and parallel readers. +/// Internal per-file cache context, shared by streams, format readers and LM wrappers. +/// It owns no readers, so builders can pass it through CacheInputStream without ownership cycles. /// Callers must stop their reader tasks before resetting or releasing a generation. -class FileReadAhead { +class PAIMON_EXPORT FileReadAhead { public: struct OpenedReaders { std::vector> readers; std::shared_ptr read_ahead; }; + static std::shared_ptr Create(const std::shared_ptr& cache, + const std::string& path); + static Result Open(const std::string& path, int64_t file_size, const ReaderBuilder* builder, const std::shared_ptr& fs, uint32_t reader_count, @@ -50,7 +59,17 @@ class FileReadAhead { const CacheConfig& cache_config, bool enable_io_metrics, const std::shared_ptr& pool); + static Result> OpenSynchronous( + const std::string& path, int64_t file_size, const ReaderBuilder* builder, + const std::shared_ptr& fs, bool cache_enabled, const CacheConfig& cache_config, + bool enable_io_metrics, WarmupLevel level, const std::shared_ptr& pool, + const std::shared_ptr& arrow_pool); + Status Warmup(PrefetchFileBatchReader* reader); + /// Prepare ranges as a format advances or LM switches projections. Inactive contexts and + /// disabled caches do nothing, including skipping the reader's range query. + Status PrepareRanges(PrefetchFileBatchReader* reader); + void SetSynchronous(bool synchronous); Status GetStatus() const; uint64_t GetErrorCount() const; void EndGeneration(); @@ -59,14 +78,15 @@ class FileReadAhead { std::shared_ptr GetMetrics() const; private: - FileReadAhead(const std::shared_ptr& cache, - const std::shared_ptr& io_metrics) - : cache_(cache), io_metrics_(io_metrics) {} + FileReadAhead(const std::shared_ptr& cache, const std::string& path) + : cache_(cache), path_(path) {} - void RegisterLateRanges(std::vector>&& ranges); + using ByteRanges = std::vector>; + Status PrepareRangesLocked(Result&& ranges, bool initial); std::shared_ptr cache_; std::shared_ptr io_metrics_; + std::string path_; mutable std::mutex mutex_; bool warmed_ = false; Status status_; 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 10bdaea18..bc29bd26f 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp @@ -35,6 +35,7 @@ #include "fmt/format.h" #include "paimon/common/predicate/predicate_filter.h" #include "paimon/common/predicate/predicate_validator.h" +#include "paimon/common/reader/file_read_ahead.h" #include "paimon/common/reader/reader_utils.h" #include "paimon/common/utils/arrow/arrow_utils.h" #include "paimon/common/utils/arrow/mem_utils.h" @@ -46,6 +47,12 @@ namespace paimon { Result> LateMaterializingFileBatchReader::Create( std::unique_ptr inner, const std::shared_ptr& arrow_pool) { + return Create(std::move(inner), arrow_pool, nullptr); +} + +Result> LateMaterializingFileBatchReader::Create( + std::unique_ptr inner, const std::shared_ptr& arrow_pool, + const std::shared_ptr& read_ahead) { if (arrow_pool == nullptr) { return Status::Invalid("arrow pool could not be nullptr."); } @@ -53,8 +60,9 @@ Result> LateMaterializingFileB return Status::Invalid("inner could not be nullptr."); } auto* prefetch_inner = dynamic_cast(inner.get()); - auto reader = std::unique_ptr( - new LateMaterializingFileBatchReader(std::move(inner), prefetch_inner, arrow_pool)); + auto reader = + std::unique_ptr(new LateMaterializingFileBatchReader( + std::move(inner), prefetch_inner, arrow_pool, read_ahead)); return reader; } @@ -73,7 +81,7 @@ Result LateMaterializingFileBatchReader::NextBatch() PAIMON_RETURN_NOT_OK( SetInnerReadSchema(payload_schema_, /*predicate=*/nullptr, matched_bitmap_)); state_ = kRunning; - PAIMON_RETURN_NOT_OK(ReportPayloadPreBufferRanges()); + PAIMON_RETURN_NOT_OK(PreparePayloadRanges()); } } @@ -263,23 +271,11 @@ 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) { +Status LateMaterializingFileBatchReader::PreparePayloadRanges() { + if (!read_ahead_ || 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(); + return read_ahead_->PrepareRanges(prefetch_inner_); } Status LateMaterializingFileBatchReader::SetReadSchema( 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 962acf6e6..b36d1202c 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader.h +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.h @@ -34,6 +34,7 @@ namespace paimon { class PredicateFilter; +class FileReadAhead; // For convenience, we abbreviate `Later Materializing` as `LatMat`. // This reader is installed below the prefetch layer (see @@ -44,6 +45,10 @@ class LateMaterializingFileBatchReader : public PrefetchFileBatchReader { static Result> Create( std::unique_ptr inner, const std::shared_ptr& arrow_pool); + static Result> Create( + std::unique_ptr inner, + const std::shared_ptr& arrow_pool, + const std::shared_ptr& read_ahead); Result NextBatch() override; @@ -93,29 +98,23 @@ class LateMaterializingFileBatchReader : public PrefetchFileBatchReader { Status SetReadRanges(const std::vector>& read_ranges) override; Result>> PreBufferRange() override { - // 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(). + // The context queries the probe ranges at warmup and the payload ranges when the + // projection changes after probing. 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, - std::shared_ptr arrow_pool) + std::shared_ptr arrow_pool, + const std::shared_ptr& read_ahead) : inner_(std::move(inner)), prefetch_inner_(prefetch_inner), - arrow_pool_(std::move(arrow_pool)) {} + arrow_pool_(std::move(arrow_pool)), + read_ahead_(read_ahead) {} /// Reset the state of the late materializing reader, does not close inner reader. void Reset(); @@ -149,10 +148,8 @@ 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(); + /// Prepare payload ranges immediately after switching the inner projection. + Status PreparePayloadRanges(); /// Returns the inner reader's prefetch interface, or an error when the format reader does not /// implement it (avro and blob do not). @@ -176,10 +173,7 @@ 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_; + std::shared_ptr read_ahead_; 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 b4295c914..cad8d5533 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 @@ -33,6 +33,9 @@ #include "arrow/array/builder_nested.h" #include "arrow/c/bridge.h" #include "gtest/gtest.h" +#include "paimon/common/io/cache_input_stream.h" +#include "paimon/common/reader/data_file_reader_factory.h" +#include "paimon/common/reader/file_read_ahead.h" #include "paimon/common/reader/late_materializing_reader_builder.h" #include "paimon/common/reader/prefetch_file_batch_reader_impl.h" #include "paimon/common/reader/reader_utils.h" @@ -841,108 +844,285 @@ 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. +// Payload preparation must reach the cache before the first payload batch is returned. +TEST_F(LateMaterializingFileBatchReaderTest, PayloadRangesPreparedBeforeFirstBatch) { using ByteRanges = std::vector>; + auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}); + auto mock = std::make_unique(data, full_type_, 3); const ByteRanges probe_ranges = {{0, 1024}}; - const ByteRanges payload_ranges = {{100000, 2048}, {200000, 4096}}; + // Keep the ranges beyond the default coalescing gap so the I/O byte assertion is exact. + const ByteRanges payload_ranges = {{100000, 2048}, {1000000, 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 cache = std::make_shared(std::make_shared(), + CacheConfig(), 0, pool_); + auto context = FileReadAhead::Create(cache, "payload-ranges"); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + LateMaterializingFileBatchReader::Create(std::move(mock), GetArrowPool(pool_), context)); auto predicate = - PredicateBuilder::GreaterOrEqual(/*field_index=*/0, "k", FieldType::BIGINT, Literal(4l)); + PredicateBuilder::GreaterOrEqual(0, "k", FieldType::BIGINT, Literal(int64_t{4})); 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); + ASSERT_EQ(up_front, probe_ranges); + ASSERT_OK(context->Warmup(reader.get())); + ASSERT_OK_AND_ASSIGN(uint64_t before, + context->GetMetrics()->GetCounter(ReadAheadCacheMetrics::IO_BYTES)); + ASSERT_EQ(before, 1024); + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch first, reader->NextBatch()); + ASSERT_FALSE(BatchReader::IsEofBatch(first)); + uint64_t first_rows = first.first->length; + ReaderUtils::ReleaseReadBatch(std::move(first)); + ASSERT_OK_AND_ASSIGN(uint64_t after, + context->GetMetrics()->GetCounter(ReadAheadCacheMetrics::IO_BYTES)); + ASSERT_EQ(after, 1024 + 2048 + 4096); + ASSERT_OK_AND_ASSIGN(std::vector rest, Collect(reader.get())); + ASSERT_EQ(first_rows + rest.size(), 6); + ASSERT_EQ(context->GetErrorCount(), 0); } -// Without a predicate there is no payload pass, so there is no late range to report. -TEST_F(LateMaterializingFileBatchReaderTest, NoCallbackReportWithoutLateMaterialization) { +TEST_F(LateMaterializingFileBatchReaderTest, NoPayloadPreparationWithoutLateMaterialization) { auto data = BuildData({0, 1, 2, 3, 4}); - auto mock = std::make_unique(data, full_type_, /*batch_size=*/2); + auto mock = std::make_unique(data, full_type_, 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)); - + auto cache = std::make_shared(std::make_shared(), + CacheConfig(), 0, pool_); + auto context = FileReadAhead::Create(cache, "plain"); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + LateMaterializingFileBatchReader::Create(std::move(mock), GetArrowPool(pool_), context)); + ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), nullptr, std::nullopt)); + ASSERT_OK(context->Warmup(reader.get())); ASSERT_OK_AND_ASSIGN(std::vector rows, Collect(reader.get())); - ASSERT_EQ(rows.size(), 5u); - EXPECT_EQ(calls, 0); + ASSERT_EQ(rows.size(), 5); + ASSERT_OK_AND_ASSIGN(uint64_t bytes, + context->GetMetrics()->GetCounter(ReadAheadCacheMetrics::IO_BYTES)); + ASSERT_EQ(bytes, 1024); } -// The probe pass matched nothing, so the payload pass never runs and reports no range. -TEST_F(LateMaterializingFileBatchReaderTest, EmptyMatchDoesNotReportPayloadRanges) { +TEST_F(LateMaterializingFileBatchReaderTest, EmptyMatchDoesNotPreparePayloadRanges) { auto data = BuildData({0, 1, 2, 3, 4}); - auto mock = std::make_unique(data, full_type_, /*batch_size=*/2); + auto mock = std::make_unique(data, full_type_, 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 cache = std::make_shared(std::make_shared(), + CacheConfig(), 0, pool_); + auto context = FileReadAhead::Create(cache, "empty-match"); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + LateMaterializingFileBatchReader::Create(std::move(mock), GetArrowPool(pool_), context)); auto predicate = - PredicateBuilder::GreaterThan(/*field_index=*/0, "k", FieldType::BIGINT, Literal(100l)); + PredicateBuilder::GreaterThan(0, "k", FieldType::BIGINT, Literal(int64_t{100})); ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate, std::nullopt)); - + ASSERT_OK(context->Warmup(reader.get())); ASSERT_OK_AND_ASSIGN(std::vector rows, Collect(reader.get())); - EXPECT_TRUE(rows.empty()); - EXPECT_EQ(calls, 0); + ASSERT_TRUE(rows.empty()); + ASSERT_OK_AND_ASSIGN(uint64_t bytes, + context->GetMetrics()->GetCounter(ReadAheadCacheMetrics::IO_BYTES)); + ASSERT_EQ(bytes, 1024); } -// 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) { +TEST_F(LateMaterializingFileBatchReaderTest, PayloadPreparationFailureIsRecordedOnce) { 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 mock = std::make_unique(data, full_type_, 2); + auto* observed = mock.get(); + auto cache = std::make_shared(std::make_shared(), + CacheConfig(), 0, pool_); + auto context = FileReadAhead::Create(cache, "payload-error"); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + LateMaterializingFileBatchReader::Create(std::move(mock), GetArrowPool(pool_), context)); auto predicate = - PredicateBuilder::GreaterOrEqual(/*field_index=*/0, "k", FieldType::BIGINT, Literal(2l)); + PredicateBuilder::GreaterOrEqual(0, "k", FieldType::BIGINT, Literal(int64_t{2})); ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate, std::nullopt)); - + ASSERT_OK(context->Warmup(reader.get())); + observed->SetPreBufferRangeStatus(Status::IOError("pre-buffer range unavailable")); ASSERT_NOK_WITH_MSG(reader->NextBatch(), "pre-buffer range unavailable"); - EXPECT_EQ(calls, 0); + ASSERT_NOK_WITH_MSG(context->PrepareRanges(observed), "pre-buffer range unavailable"); + ASSERT_EQ(context->GetErrorCount(), 1); + context->Reset(); + observed->SetPreBufferRangeStatus(Status::OK()); + ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate, std::nullopt)); + ASSERT_OK(context->Warmup(reader.get())); + ASSERT_OK_AND_ASSIGN(std::vector rows, Collect(reader.get())); + ASSERT_EQ(rows.size(), 3); + ASSERT_OK(context->GetStatus()); +} + +TEST_F(LateMaterializingFileBatchReaderTest, IncrementalReaderUsesSharedStreamContext) { + class IncrementalReader : public MockFileBatchReader { + public: + IncrementalReader(const std::shared_ptr& data, + const std::shared_ptr& context) + : MockFileBatchReader(data, data->type(), 2), context_(context) { + EnableRandomizeBatchSize(false); + } + Result>> PreBufferRange() override { + return std::vector>{{offset_, 4}}; + } + Result NextBatchWithBitmap() override { + if (++reads_ == 2) { + offset_ = 100; + PAIMON_RETURN_NOT_OK(context_->PrepareRanges(this)); + } + return MockFileBatchReader::NextBatchWithBitmap(); + } + + private: + std::shared_ptr context_; + uint64_t offset_ = 0; + uint32_t reads_ = 0; + }; + class Builder : public MockFormatReaderBuilder { + public: + explicit Builder(const std::shared_ptr& data) + : MockFormatReaderBuilder(data, data->type(), 2), data_(data) {} + Result> Build( + const std::shared_ptr& input) const override { + auto stream = std::dynamic_pointer_cast(input); + if (!stream || !stream->GetReadAhead()) { + return Status::Invalid("missing shared read-ahead context"); + } + return std::make_unique(data_, stream->GetReadAhead()); + } + + private: + std::shared_ptr data_; + }; + auto data = BuildData({0, 1, 2, 3, 4, 5}); + LateMaterializingReaderBuilder builder(std::make_unique(data), GetArrowPool(pool_)); + ASSERT_OK_AND_ASSIGN( + FileReadAhead::OpenedReaders opened, + FileReadAhead::Open("incremental", 0, &builder, std::make_shared(), 1, + nullptr, true, CacheConfig(), false, pool_)); + auto reader = opened.readers[0]; + ArrowSchema schema; + ASSERT_TRUE(arrow::ExportSchema(*arrow::schema(full_fields_), &schema).ok()); + ASSERT_OK(reader->SetReadSchema(&schema, nullptr, std::nullopt)); + ASSERT_OK(opened.read_ahead->Warmup(reader.get())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr rows, CollectStruct(reader.get())); + ASSERT_EQ(rows->length(), 6); + ASSERT_OK_AND_ASSIGN(uint64_t bytes, opened.read_ahead->GetMetrics()->GetCounter( + ReadAheadCacheMetrics::IO_BYTES)); + ASSERT_EQ(bytes, 8); + opened.read_ahead->ReleaseBuffers(); +} + +TEST_F(LateMaterializingFileBatchReaderTest, CacheDisabledSkipsPayloadRangeQueries) { + class FailingRangeReader : public MockFileBatchReader { + public: + FailingRangeReader(const std::shared_ptr& data, uint32_t* calls) + : MockFileBatchReader(data, data->type(), 2), calls_(calls) { + EnableRandomizeBatchSize(false); + } + Result>> PreBufferRange() override { + ++*calls_; + return Status::IOError("range query must not run with caching disabled"); + } + + private: + uint32_t* calls_; + }; + class Builder : public MockFormatReaderBuilder { + public: + Builder(const std::shared_ptr& data, uint32_t* calls) + : MockFormatReaderBuilder(data, data->type(), 2), data_(data), calls_(calls) {} + Result> Build( + const std::shared_ptr& input) const override { + return std::make_unique(data_, calls_); + } + + private: + std::shared_ptr data_; + uint32_t* calls_; + }; + auto data = BuildData({0, 1, 2, 3, 4, 5}); + auto fs = std::make_shared(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr executor, CreateDefaultExecutor(1)); + for (WarmupLevel level : {WarmupLevel::NONE, WarmupLevel::RAW}) { + for (int32_t mode = 0; mode < 3; ++mode) { + SCOPED_TRACE(mode); + uint32_t calls = 0; + LateMaterializingReaderBuilder builder(std::make_unique(data, &calls), + GetArrowPool(pool_)); + DataFileReadOptions options; + options.read_ahead_cache_enabled = false; + options.warmup_level = level; + options.prefetch_enabled = mode != 0; + options.adaptive_prefetch_strategy = mode == 2; + options.prefetch_max_parallel_num = 1; + options.read_batch_size = 1; + options.prefetch_batch_count = 1; + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + DataFileReaderFactory::Open("parquet", "cache-disabled", 0, &builder, options, fs, + executor, pool_, GetArrowPool(pool_))); + auto predicate = + PredicateBuilder::GreaterOrEqual(0, "k", FieldType::BIGINT, Literal(int64_t{2})); + ArrowSchema c_schema; + ASSERT_TRUE(arrow::ExportSchema(*arrow::schema(full_fields_), &c_schema).ok()); + ASSERT_OK(reader->SetReadSchema(&c_schema, predicate, std::nullopt)); + reader->Warmup(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, + CollectStruct(reader.get())); + ASSERT_EQ(result->length(), 4); + ASSERT_EQ(calls, 0); + reader->Close(); + } + } +} + +TEST_F(LateMaterializingFileBatchReaderTest, SharedContextFailureIsNotCountedAsDecodeFailure) { + class Reader : public MockFileBatchReader { + public: + explicit Reader(const std::shared_ptr& data) + : MockFileBatchReader(data, data->type(), 2) {} + Result>> PreBufferRange() override { + if (++queries_ > 1) { + return Status::IOError("payload preparation failed"); + } + return std::vector>(); + } + + private: + uint32_t queries_ = 0; + }; + class Builder : public MockFormatReaderBuilder { + public: + explicit Builder(const std::shared_ptr& data) + : MockFormatReaderBuilder(data, data->type(), 2), data_(data) {} + Result> Build( + const std::shared_ptr& input) const override { + return std::make_unique(data_); + } + + private: + std::shared_ptr data_; + }; + auto data = BuildData({0, 1, 2, 3, 4, 5}); + LateMaterializingReaderBuilder builder(std::make_unique(data), GetArrowPool(pool_)); + DataFileReadOptions options; + options.prefetch_enabled = true; + options.prefetch_max_parallel_num = 1; + options.read_batch_size = 2; + options.prefetch_batch_count = 2; + ASSERT_OK_AND_ASSIGN(std::shared_ptr executor, CreateDefaultExecutor(1)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, + DataFileReaderFactory::Open("parquet", "payload-error", 0, &builder, + options, std::make_shared(), + executor, pool_, GetArrowPool(pool_))); + ArrowSchema schema; + ASSERT_TRUE(arrow::ExportSchema(*arrow::schema(full_fields_), &schema).ok()); + auto predicate = + PredicateBuilder::GreaterOrEqual(0, "k", FieldType::BIGINT, Literal(int64_t{2})); + ASSERT_OK(reader->SetReadSchema(&schema, predicate, std::nullopt)); + reader->Warmup(); + ASSERT_NOK_WITH_MSG(CollectStruct(reader.get()), "payload preparation failed"); + reader->Close(); + ASSERT_OK_AND_ASSIGN(uint64_t errors, + reader->GetReaderMetrics()->GetCounter(PrefetchMetrics::ERRORS)); + ASSERT_EQ(errors, 1); } -// End to end below the prefetch layer: the ranges the inner LM reader reports mid-read must reach +// End to end below the prefetch layer: the ranges the inner LM reader prepares 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}); diff --git a/src/paimon/common/reader/late_materializing_reader_builder.h b/src/paimon/common/reader/late_materializing_reader_builder.h index 56571fcb9..c23c220bd 100644 --- a/src/paimon/common/reader/late_materializing_reader_builder.h +++ b/src/paimon/common/reader/late_materializing_reader_builder.h @@ -22,6 +22,7 @@ #include #include +#include "paimon/common/io/cache_input_stream.h" #include "paimon/common/reader/late_materializing_file_batch_reader.h" #include "paimon/format/reader_builder.h" #include "paimon/memory/memory_pool.h" @@ -55,9 +56,11 @@ class LateMaterializingReaderBuilder : public ReaderBuilder { const std::shared_ptr& stream) const override { PAIMON_ASSIGN_OR_RAISE(std::unique_ptr format_reader, inner_->Build(stream)); - PAIMON_ASSIGN_OR_RAISE( - std::unique_ptr reader, - LateMaterializingFileBatchReader::Create(std::move(format_reader), arrow_pool_)); + auto cache_stream = std::dynamic_pointer_cast(stream); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader, + LateMaterializingFileBatchReader::Create( + std::move(format_reader), arrow_pool_, + cache_stream ? cache_stream->GetReadAhead() : nullptr)); return std::unique_ptr(std::move(reader)); } 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 9e7ffa0bb..93ddd1cc3 100644 --- a/src/paimon/common/reader/prefetch_file_batch_reader_impl.cpp +++ b/src/paimon/common/reader/prefetch_file_batch_reader_impl.cpp @@ -169,7 +169,6 @@ PrefetchFileBatchReaderImpl::PrefetchFileBatchReaderImpl( PrefetchFileBatchReaderImpl::~PrefetchFileBatchReaderImpl() { (void)CleanUp(); - ClearPreBufferRangeCallbacks(); } Status PrefetchFileBatchReaderImpl::SetReadSchema( @@ -214,6 +213,7 @@ Status PrefetchFileBatchReaderImpl::RefreshReadRangesAfterCleanUp() { prefetch_metrics_->adaptive_disabled_count.fetch_add(1, kMetricsMemoryOrder); } need_prefetch_ = need_prefetch; + read_ahead_->SetSynchronous(!need_prefetch_); prefetch_metrics_->enabled.store(need_prefetch_, kMetricsMemoryOrder); prefetch_metrics_->read_ranges_total.fetch_add(read_ranges.size(), kMetricsMemoryOrder); std::vector> filtered_ranges = @@ -577,14 +577,6 @@ void PrefetchFileBatchReaderImpl::WarmCacheOnce() { (void)read_ahead_->Warmup(readers_[0].get()); } -void PrefetchFileBatchReaderImpl::ClearPreBufferRangeCallbacks() { - // CleanUp() has joined the background thread. Detach the weak state callbacks so retained - // sub-readers cannot report more byte ranges after this reader releases its generation. - for (const auto& sub_reader : readers_) { - sub_reader->SetPreBufferRangeCallback(nullptr); - } -} - void PrefetchFileBatchReaderImpl::Warmup() { if (warmup_level_ != WarmupLevel::RAW) { return; @@ -763,7 +755,7 @@ Result PrefetchFileBatchReaderImpl::GetNextRowToRead() const { } void PrefetchFileBatchReaderImpl::SetReadStatus(const Status& status) { - if (!status.ok()) { + if (!status.ok() && status != read_ahead_->GetStatus()) { prefetch_metrics_->errors.fetch_add(1, kMetricsMemoryOrder); } std::unique_lock lock(rw_mutex_); @@ -795,7 +787,6 @@ void PrefetchFileBatchReaderImpl::Close() { for (size_t i = 1; i < readers_.size(); ++i) { readers_[i]->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 b1e4a7af2..c7de0c126 100644 --- a/src/paimon/common/reader/prefetch_file_batch_reader_impl.h +++ b/src/paimon/common/reader/prefetch_file_batch_reader_impl.h @@ -137,9 +137,6 @@ class PrefetchFileBatchReaderImpl : public PrefetchFileBatchReader { /// Shares byte-range registration with the synchronous warmup path. void WarmCacheOnce(); - /// 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; diff --git a/src/paimon/common/reader/warmup_file_batch_reader.cpp b/src/paimon/common/reader/warmup_file_batch_reader.cpp index 183edcd49..41277180e 100644 --- a/src/paimon/common/reader/warmup_file_batch_reader.cpp +++ b/src/paimon/common/reader/warmup_file_batch_reader.cpp @@ -35,7 +35,6 @@ WarmupFileBatchReader::WarmupFileBatchReader(std::shared_ptrSetPreBufferRangeCallback(nullptr); read_ahead_->ReleaseBuffers(); } @@ -135,7 +134,6 @@ void WarmupFileBatchReader::Close() { } ResetReadState(); closed_ = true; - reader_->SetPreBufferRangeCallback(nullptr); read_ahead_->ReleaseBuffers(); reader_->Close(); } diff --git a/src/paimon/common/reader/warmup_file_batch_reader_test.cpp b/src/paimon/common/reader/warmup_file_batch_reader_test.cpp index 80d8c4b9f..717b0a108 100644 --- a/src/paimon/common/reader/warmup_file_batch_reader_test.cpp +++ b/src/paimon/common/reader/warmup_file_batch_reader_test.cpp @@ -27,6 +27,8 @@ #include "arrow/c/bridge.h" #include "arrow/ipc/json_simple.h" #include "gtest/gtest.h" +#include "paimon/common/io/cache_input_stream.h" +#include "paimon/common/reader/data_file_reader_factory.h" #include "paimon/common/reader/reader_utils.h" #include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/checked_cast.h" @@ -48,6 +50,9 @@ class ObservedReader : public MockFileBatchReader { const std::shared_ptr& stream) : MockFileBatchReader(data, data->type(), /*read_batch_size=*/2), stream_(stream) { EnableRandomizeBatchSize(false); + if (auto cache_stream = std::dynamic_pointer_cast(stream)) { + read_ahead_ = cache_stream->GetReadAhead(); + } } Result NextBatchWithBitmap() override { @@ -58,29 +63,30 @@ class ObservedReader : public MockFileBatchReader { Result>> PreBufferRange() override { PAIMON_RETURN_NOT_OK(range_status); - return std::vector>{{0, 4}}; + return std::vector>{{range_offset_, range_length}}; } - void SetPreBufferRangeCallback(PreBufferRangeCallback callback) override { - callback_ = std::move(callback); - } - - void ReportLateRanges(uint64_t offset) { - if (callback_) { - callback_({{offset, 4}}); - } + Status PrepareLateRanges(uint64_t offset) { + uint64_t previous_offset = range_offset_; + range_offset_ = offset; + Status status = read_ahead_ ? read_ahead_->PrepareRanges(this) : Status::OK(); + range_offset_ = previous_offset; + return status; } uint32_t read_count = 0; std::thread::id last_read_thread; bool read_bytes = false; + uint64_t range_length = 4; + uint64_t read_stride = 0; Status range_status; private: Result Read() { if (read_bytes) { char bytes[4]; - PAIMON_ASSIGN_OR_RAISE(int64_t size, stream_->Read(bytes, 4, 0)); + PAIMON_ASSIGN_OR_RAISE(int64_t size, + stream_->Read(bytes, 4, (read_count - 1) * read_stride)); if (size != 4) { return Status::IOError("short observed read"); } @@ -89,7 +95,8 @@ class ObservedReader : public MockFileBatchReader { } std::shared_ptr stream_; - PreBufferRangeCallback callback_; + std::shared_ptr read_ahead_; + uint64_t range_offset_ = 0; }; class ObservedBuilder : public ReaderBuilder { @@ -102,9 +109,13 @@ class ObservedBuilder : public ReaderBuilder { Result> Build( const std::shared_ptr& stream) const override { - return std::make_unique(data_, stream); + auto reader = std::make_unique(data_, stream); + last_reader = reader.get(); + return reader; } + mutable ObservedReader* last_reader = nullptr; + private: std::shared_ptr data_; }; @@ -314,23 +325,90 @@ TEST_F(WarmupFileBatchReaderTest, RawRangeAndAsyncFailuresAreDeferredUntilRead) ASSERT_NOK_WITH_MSG(reader->NextBatchWithBitmap(), "async warmup failed"); } +TEST_F(WarmupFileBatchReaderTest, SynchronousAndAdaptiveFallbackBoundRetainedBytes) { + constexpr uint64_t kMiB = 1024 * 1024; + for (bool adaptive_fallback : {false, true}) { + for (WarmupLevel level : {WarmupLevel::NONE, WarmupLevel::RAW}) { + auto pool = std::shared_ptr(GetMemoryPool()); + ObservedBuilder builder(data_); + DataFileReadOptions options; + options.warmup_level = level; + options.prefetch_enabled = adaptive_fallback; + options.adaptive_prefetch_strategy = true; + options.prefetch_max_parallel_num = 1; + options.prefetch_batch_count = 1; + options.read_batch_size = 1; + options.cache_config.SetRangeSizeLimit(kMiB); + options.cache_config.SetHoleSizeLimit(0); + options.cache_config.SetPreBufferLimit(2 * kMiB); + options.cache_config.SetBlockSize(0); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + DataFileReaderFactory::Open("parquet", "bounded-sync", 32 * kMiB, &builder, options, + std::make_shared(), executor_, pool, + GetArrowPool(pool))); + ASSERT_NE(builder.last_reader, nullptr); + builder.last_reader->range_length = 32 * kMiB; + builder.last_reader->read_stride = 8 * kMiB; + builder.last_reader->read_bytes = true; + ASSERT_OK(Configure(reader.get())); + reader->Warmup(); + uint64_t row_count = 0; + while (true) { + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatchWithBitmap batch, + reader->NextBatchWithBitmap()); + if (BatchReader::IsEofBatch(batch)) { + break; + } + row_count += batch.first.first->length; + ReaderUtils::ReleaseReadBatch(std::move(batch.first)); + } + ASSERT_EQ(row_count, 6); + ASSERT_LE(pool->MaxMemoryUsage(), 2 * kMiB); + reader->Close(); + ASSERT_EQ(pool->CurrentUsage(), 0); + } + } +} + +TEST_F(WarmupFileBatchReaderTest, InvalidInitialAndLateRangesReportFileOnRead) { + for (bool late : {false, true}) { + ObservedBuilder builder(data_); + ASSERT_OK_AND_ASSIGN(FileReadAhead::OpenedReaders opened, + FileReadAhead::Open("range-bounds-file", late ? 4 : 2, &builder, + std::make_shared(), 1, nullptr, + true, CacheConfig(), false, pool_)); + auto reader = std::make_unique( + opened.readers[0], opened.read_ahead, WarmupLevel::RAW, GetArrowPool(pool_)); + ASSERT_OK(Configure(reader.get())); + reader->Warmup(); + if (late) { + ASSERT_NOK_WITH_MSG(builder.last_reader->PrepareLateRanges(2), "exceeds file size"); + } + ASSERT_NOK_WITH_MSG(reader->NextBatchWithBitmap(), "range-bounds-file"); + ASSERT_NOK_WITH_MSG(reader->NextBatchWithBitmap(), "exceeds file size"); + ASSERT_EQ(builder.last_reader->read_count, 0); + reader->Close(); + } +} + TEST_F(WarmupFileBatchReaderTest, LateRangesAreRegisteredOnlyInActiveGeneration) { ASSERT_OK_AND_ASSIGN(auto reader, Open(WarmupLevel::RAW, true)); ASSERT_OK(Configure(reader.get())); - observed_->ReportLateRanges(16); + ASSERT_OK(observed_->PrepareLateRanges(16)); ASSERT_EQ(Counter(reader.get(), ReadAheadCacheMetrics::IO_COUNT), 0); reader->Warmup(); - observed_->ReportLateRanges(16); + ASSERT_OK(observed_->PrepareLateRanges(16)); ASSERT_EQ(Counter(reader.get(), ReadAheadCacheMetrics::IO_COUNT), 2); - observed_->ReportLateRanges(16); + ASSERT_OK(observed_->PrepareLateRanges(16)); ASSERT_EQ(Counter(reader.get(), ReadAheadCacheMetrics::IO_COUNT), 2); ASSERT_OK(Configure(reader.get())); - observed_->ReportLateRanges(32); + ASSERT_OK(observed_->PrepareLateRanges(32)); ASSERT_EQ(Counter(reader.get(), ReadAheadCacheMetrics::IO_COUNT), 0); reader->Warmup(); ASSERT_EQ(Counter(reader.get(), ReadAheadCacheMetrics::IO_COUNT), 1); reader->Close(); - observed_->ReportLateRanges(32); + ASSERT_OK(observed_->PrepareLateRanges(32)); ASSERT_EQ(Counter(reader.get(), ReadAheadCacheMetrics::IO_COUNT), 1); } diff --git a/src/paimon/common/utils/read_ahead_cache.cpp b/src/paimon/common/utils/read_ahead_cache.cpp index 68dc18034..2633b2a29 100644 --- a/src/paimon/common/utils/read_ahead_cache.cpp +++ b/src/paimon/common/utils/read_ahead_cache.cpp @@ -23,13 +23,16 @@ #include "paimon/common/utils/read_ahead_cache.h" #include +#include #include +#include #include #include #include #include #include +#include "fmt/format.h" #include "paimon/common/memory/bytes_utils.h" #include "paimon/common/metrics/atomic_counter_pair.h" #include "paimon/common/utils/byte_range_combiner.h" @@ -49,6 +52,10 @@ struct PendingFetch { std::shared_ptr> promise; }; +struct ReadAheadCacheBudget { + std::atomic resident_bytes{0}; +}; + namespace { // Copy the requested window out of the covering entries into dest. The @@ -69,6 +76,17 @@ void CopyRangeFromEntries(const std::vector& covering, const By } // namespace Result> ReadAheadCache::AddRanges(std::vector&& ranges) { + // Check before coalescing splits ranges and allocates entries. Zero is the existing + // unknown-size convention; integer bounds are still checked by the combiner in that case. + if (file_size_ > 0) { + for (const auto& range : ranges) { + if (range.offset > file_size_ || range.length > file_size_ - range.offset) { + return Status::Invalid( + fmt::format("byte range (offset={}, length={}) exceeds file size {}", + range.offset, range.length, file_size_)); + } + } + } // 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, @@ -188,6 +206,18 @@ void ReadAheadCache::PreBuffer(uint64_t offset) { } size_t total_bytes = 0; + auto window_end = it; + while (bounded_retention_ && window_end != ranges_.end() && + window_end->range.length <= config_.GetPreBufferLimit() - total_bytes) { + total_bytes += window_end->range.length; + ++window_end; + } + const uint64_t window_begin_offset = it->range.offset; + const uint64_t window_end_offset = + window_end == it + ? window_begin_offset + : std::prev(window_end)->range.offset + std::prev(window_end)->range.length; + total_bytes = 0; for (auto entry = it; entry != ranges_.end(); ++entry) { total_bytes += entry->range.length; if (total_bytes > config_.GetPreBufferLimit()) { @@ -196,14 +226,18 @@ void ReadAheadCache::PreBuffer(uint64_t offset) { if (entry->Published()) { continue; } + if (bounded_retention_ && + !MakeRoomLocked(entry->range.length, window_begin_offset, window_end_offset)) { + break; + } auto promise = std::make_shared>(); auto future = promise->get_future(); - auto buffer = AllocateBytesKeepingPoolAlive(entry->range.length, memory_pool_); + auto buffer = AllocateRangeBuffer(entry->range.length); 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. + // issuing a second fetch for the same bytes. In-flight entries are never evicted, + // so every pending fetch remains reachable for Close()/Reset() to wait on. entry->buffer = std::move(buffer); entry->future = std::move(future).share(); } @@ -213,10 +247,58 @@ void ReadAheadCache::PreBuffer(uint64_t offset) { DispatchFetches(fetches); } +void ReadAheadCache::SetBoundedRetention(bool enabled) { + std::unique_lock lock(rw_mutex_); + bounded_retention_ = enabled; +} + +bool ReadAheadCache::MakeRoomLocked(uint64_t size, uint64_t window_begin, uint64_t window_end) { + const uint64_t limit = config_.GetPreBufferLimit(); + if (size > limit) { + return false; + } + for (size_t visited = 0; visited < ranges_.size(); ++visited) { + if (budget_->resident_bytes.load(std::memory_order_acquire) <= limit - size) { + return true; + } + auto& entry = ranges_[eviction_cursor_ % ranges_.size()]; + eviction_cursor_ = (eviction_cursor_ + 1) % ranges_.size(); + // Keep this window, pending I/O, failed results and buffers being copied or retained by + // a filesystem callback. Re-reading an evicted range simply publishes it again. + if (!entry.Published() || + (entry.range.offset >= window_begin && entry.range.offset < window_end) || + entry.buffer.use_count() != 1 || + entry.future.wait_for(std::chrono::seconds(0)) != std::future_status::ready || + !entry.future.get().ok()) { + continue; + } + entry.buffer.reset(); + entry.future = {}; + } + return budget_->resident_bytes.load(std::memory_order_acquire) <= limit - size; +} + +std::shared_ptr ReadAheadCache::AllocateRangeBuffer(uint64_t size) { + auto buffer = AllocateBytesKeepingPoolAlive(size, memory_pool_); + budget_->resident_bytes.fetch_add(size, std::memory_order_relaxed); + // Account until the LAST owner releases the buffer, even if a callback retains it after + // Reset(). The captured buffer also retains the allocating memory pool. + return std::shared_ptr( + buffer.get(), [buffer, budget = budget_, size](Bytes* bytes) mutable { + assert(bytes == buffer.get()); + buffer.reset(); + budget->resident_bytes.fetch_sub(size, std::memory_order_release); + }); +} + 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) { + : stream_(stream), + config_(config), + file_size_(file_size), + budget_(std::make_shared()), + 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. if (file_size > 0 && config_.GetBlockSize() > 0 && config_.GetBlockCacheLimit() > 0) { @@ -262,11 +344,12 @@ void ReadAheadCache::ReleasePrefetchBuffers() { // 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 + // can observe or mutate them. In-flight entries are never evicted, so waiting on these // 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_); + eviction_cursor_ = 0; } // Wait OUTSIDE the lock; the buffers are released when to_release destructs. for (auto& entry : to_release) { diff --git a/src/paimon/common/utils/read_ahead_cache.h b/src/paimon/common/utils/read_ahead_cache.h index 4eec73836..29977ddaa 100644 --- a/src/paimon/common/utils/read_ahead_cache.h +++ b/src/paimon/common/utils/read_ahead_cache.h @@ -44,6 +44,7 @@ class Metrics; class FileBlockCache; class Bytes; struct PendingFetch; +struct ReadAheadCacheBudget; /// Metric names for the read-ahead cache. class PAIMON_EXPORT ReadAheadCacheMetrics { @@ -137,9 +138,9 @@ struct RangeCacheEntry { /// 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 -/// a single data file, whose size is bounded by the reader's scan scope. +/// By default published ranges stay cached until ReleaseBuffers() or Reset(). Synchronous +/// readers enable bounded retention: completed, unpinned buffers outside the current fetch +/// window can be evicted, and the pre-buffer limit also bounds resident range buffers. /// /// Reads that the prefetched ranges do not cover - a reader reads the metadata /// of its file before any range is registered - are served by a FileBlockCache @@ -200,6 +201,11 @@ class PAIMON_EXPORT ReadAheadCache { /// @param from_offset The offset to start fetching from, typically one AddRanges() returned. void Warmup(uint64_t from_offset); + /// Configure while reader tasks are stopped. Synchronous reads bound retained range buffers + /// by the pre-buffer limit; the block cache keeps its separate budget. In-flight reads and + /// buffers retained by callbacks count against the range budget, including across Reset(). + void SetBoundedRetention(bool enabled); + /// 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 @@ -254,6 +260,8 @@ class PAIMON_EXPORT ReadAheadCache { /// 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); + bool MakeRoomLocked(uint64_t size, uint64_t window_begin, uint64_t window_end); + std::shared_ptr AllocateRangeBuffer(uint64_t size); /// Clear the prefetch state, waiting for the fetches still writing into the /// entry buffers. Leaves the block cache untouched. @@ -261,10 +269,14 @@ class PAIMON_EXPORT ReadAheadCache { std::shared_ptr stream_; CacheConfig config_; + uint64_t file_size_; + bool bounded_retention_ = false; + std::shared_ptr budget_; // 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_; + size_t eviction_cursor_ = 0; std::shared_ptr memory_pool_; mutable std::shared_mutex rw_mutex_; // Caches the reads that no registered range covers, or null when the block diff --git a/src/paimon/common/utils/read_ahead_cache_test.cpp b/src/paimon/common/utils/read_ahead_cache_test.cpp index c9f7454a1..4439712f6 100644 --- a/src/paimon/common/utils/read_ahead_cache_test.cpp +++ b/src/paimon/common/utils/read_ahead_cache_test.cpp @@ -21,6 +21,7 @@ #include #include +#include #include #include #include @@ -213,6 +214,146 @@ TEST(TestReadAheadCache, TestNoEvictionKeepsAllRanges) { AssertReadEquals({0, 5}, "abcde", &cache); } +TEST(TestReadAheadCache, BoundedRetentionLimitsPeakMemoryAndAllowsRewind) { + constexpr uint64_t kMiB = 1024 * 1024; + std::string content(32 * kMiB, 'a'); + std::unique_ptr dir; + auto stream = OpenTestFile(&dir, "bounded", content); + auto pool = TestPool(); + CacheConfig config = TestCacheConfig(kMiB, 0, 2 * kMiB); + config.SetBlockSize(0); + ReadAheadCache cache(stream, config, content.size(), pool); + cache.SetBoundedRetention(true); + ASSERT_OK(cache.AddRanges({{0, 16 * kMiB}})); + cache.Warmup(); + ASSERT_EQ(pool->CurrentUsage(), 2 * kMiB); + for (uint64_t i = 0; i < 32; ++i) { + if (i == 16) { + ASSERT_OK(cache.AddRanges({{16 * kMiB, 16 * kMiB}})); + } + AssertReadEquals({i * kMiB, kMiB}, std::string(kMiB, 'a'), &cache); + ASSERT_LE(pool->CurrentUsage(), 2 * kMiB); + } + AssertReadEquals({0, kMiB}, std::string(kMiB, 'a'), &cache); + ASSERT_LE(pool->MaxMemoryUsage(), 2 * kMiB); + cache.ReleaseBuffers(); + ASSERT_EQ(pool->CurrentUsage(), 0); +} + +TEST(TestReadAheadCache, BoundedRetentionCountsPendingAndRetainedCallbacksAcrossReset) { + std::unique_ptr dir; + auto stream = std::make_shared( + OpenTestFile(&dir, "pending", "abcdefghijklmnopqrstuvwxyz")); + auto pool = TestPool(); + CacheConfig config = TestCacheConfig(4, 0, 8); + config.SetBlockSize(0); + ReadAheadCache cache(stream, config, 26, pool); + cache.SetBoundedRetention(true); + ASSERT_OK(cache.AddRanges({{0, 24}})); + cache.Warmup(); + ASSERT_EQ(stream->AsyncReadCount(), 2); + AssertReadMiss({8, 4}, &cache); + ASSERT_EQ(stream->AsyncReadCount(), 2); + stream->ReleaseAllKeepingCallbacks(); + AssertReadMiss({8, 4}, &cache); + ASSERT_EQ(pool->CurrentUsage(), 8); + cache.Reset(); + ASSERT_OK(cache.AddRanges({{12, 4}})); + cache.Warmup(); + ASSERT_EQ(stream->AsyncReadCount(), 2); + stream->DropCompletedCallbacks(); + ASSERT_EQ(pool->CurrentUsage(), 0); + cache.Warmup(); + ASSERT_EQ(stream->AsyncReadCount(), 3); + stream->ReleaseAll(); + AssertReadEquals({12, 4}, "mnop", &cache); + ASSERT_LE(pool->MaxMemoryUsage(), 8); + cache.ReleaseBuffers(); + ASSERT_EQ(pool->CurrentUsage(), 0); +} + +TEST(TestReadAheadCache, BoundedRetentionCountsBuffersFromEarlierParallelGeneration) { + std::unique_ptr dir; + auto stream = std::make_shared( + OpenTestFile(&dir, "mode-switch", "abcdefghijklmnop")); + auto pool = TestPool(); + CacheConfig config = TestCacheConfig(4, 0, 8); + config.SetBlockSize(0); + ReadAheadCache cache(stream, config, 16, pool); + ASSERT_OK(cache.AddRanges({{0, 8}})); + cache.Warmup(); + stream->ReleaseAllKeepingCallbacks(); + cache.Reset(); + cache.SetBoundedRetention(true); + ASSERT_OK(cache.AddRanges({{8, 4}})); + cache.Warmup(); + ASSERT_EQ(stream->AsyncReadCount(), 2); + ASSERT_EQ(pool->CurrentUsage(), 8); + stream->DropCompletedCallbacks(); + cache.Warmup(); + ASSERT_EQ(stream->AsyncReadCount(), 3); + stream->ReleaseAll(); + AssertReadEquals({8, 4}, "ijkl", &cache); + ASSERT_LE(pool->MaxMemoryUsage(), 8); +} + +TEST(TestReadAheadCache, BoundedRetentionDeclinesOversizedReadsWithoutAllocating) { + std::unique_ptr dir; + auto stream = OpenTestFile(&dir, "oversized", "abcdefghijklmnop"); + auto pool = TestPool(); + CacheConfig config = TestCacheConfig(8, 0, 4); + config.SetBlockSize(0); + ReadAheadCache cache(stream, config, 16, pool); + cache.SetBoundedRetention(true); + ASSERT_OK(cache.AddRanges({{0, 16}})); + cache.Warmup(); + AssertReadMiss({0, 8}, &cache); + ASSERT_EQ(pool->MaxMemoryUsage(), 0); + char bytes[8]; + ASSERT_OK_AND_ASSIGN(int64_t read, stream->Read(bytes, 8, 0)); + ASSERT_EQ(read, 8); + ASSERT_EQ(std::string(bytes, 8), "abcdefgh"); +} + +TEST(TestReadAheadCache, BoundedRetentionKeepsFailedFetchErrors) { + std::unique_ptr dir; + auto stream = + std::make_shared(OpenTestFile(&dir, "failed", "abcdefghijklmnop")); + CacheConfig config = TestCacheConfig(4, 0, 4); + config.SetBlockSize(0); + ReadAheadCache cache(stream, config, 16, TestPool()); + cache.SetBoundedRetention(true); + ASSERT_OK(cache.AddRanges({{0, 16}})); + cache.Warmup(); + stream->FailAll(Status::IOError("failed fetch")); + AssertReadMiss({4, 4}, &cache); + char buffer[4]; + ASSERT_NOK_WITH_MSG(cache.Read({0, 4}, buffer), "failed fetch"); + ASSERT_EQ(stream->AsyncReadCount(), 1); +} + +TEST(TestReadAheadCache, RejectOutOfFileRangesBeforeSplitting) { + std::unique_ptr dir; + auto stream = OpenTestFile(&dir, "bounds", std::string(1024, 'a')); + CacheConfig config = TestCacheConfig(1024, 0, 2048); + config.SetBlockSize(0); + ReadAheadCache cache(stream, config, 1024, TestPool()); + for (const ByteRange& range : + std::vector{{512, 32 * 1024 * 1024}, + {1025, 0}, + {1024, 1}, + {std::numeric_limits::max(), 1}}) { + ASSERT_NOK_WITH_MSG(cache.AddRanges({{0, 1}, range}), "exceeds file size 1024"); + ASSERT_TRUE(cache.ranges_.empty()); + } + ASSERT_OK(cache.AddRanges({{1024, 0}, {1023, 1}})); + ASSERT_EQ(cache.ranges_.size(), 1); + ReadAheadCache unknown(stream, TestCacheConfig(1024 * 1024, 0, 8), 0, TestPool()); + ASSERT_OK(unknown.AddRanges({{512, 32 * 1024 * 1024}})); + ASSERT_NOK_WITH_MSG(unknown.AddRanges({{1, std::numeric_limits::max()}}), + "int64 bound"); +} + // Test that Read() hits and misses are recorded in the cache metrics. TEST(TestReadAheadCache, TestMetrics) { CacheConfig config = diff --git a/src/paimon/format/orc/orc_file_batch_reader.cpp b/src/paimon/format/orc/orc_file_batch_reader.cpp index 526ff7cd2..91fb931ab 100644 --- a/src/paimon/format/orc/orc_file_batch_reader.cpp +++ b/src/paimon/format/orc/orc_file_batch_reader.cpp @@ -88,16 +88,26 @@ void CollectAllColumnIds(const ::orc::Type* type, std::vector* column_ OrcFileBatchReader::OrcFileBatchReader(std::unique_ptr<::orc::ReaderMetrics>&& reader_metrics, std::unique_ptr&& reader, - const std::map& options) + const std::map& options, + const std::shared_ptr& read_ahead) : options_(options), reader_metrics_(std::move(reader_metrics)), reader_(std::move(reader)), - metrics_(std::make_shared()) {} + metrics_(std::make_shared()), + read_ahead_(read_ahead) {} Result> OrcFileBatchReader::Create( std::unique_ptr<::orc::InputStream>&& input_stream, const std::shared_ptr& read_memory, const std::map& options, int32_t batch_size) { + return Create(std::move(input_stream), read_memory, options, batch_size, nullptr); +} + +Result> OrcFileBatchReader::Create( + std::unique_ptr<::orc::InputStream>&& input_stream, + const std::shared_ptr& read_memory, + const std::map& options, int32_t batch_size, + const std::shared_ptr& read_ahead) { assert(input_stream); std::string file_name = input_stream->getName(); try { @@ -126,8 +136,8 @@ Result> OrcFileBatchReader::Create( PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader_wrapper, OrcReaderWrapper::Create(std::move(reader), file_name, batch_size, natural_read_size, options, read_memory)); - auto orc_file_batch_reader = std::unique_ptr( - new OrcFileBatchReader(std::move(reader_metrics), std::move(reader_wrapper), options)); + auto orc_file_batch_reader = std::unique_ptr(new OrcFileBatchReader( + std::move(reader_metrics), std::move(reader_wrapper), options, read_ahead)); PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<::ArrowSchema> file_schema, orc_file_batch_reader->GetFileSchema()); PAIMON_RETURN_NOT_OK(orc_file_batch_reader->SetReadSchema( @@ -203,6 +213,7 @@ Status OrcFileBatchReader::SetReadSchema(::ArrowSchema* read_schema, target_column_ids_ = target_column_ids; previous_batch_row_count_ = 0; PAIMON_RETURN_NOT_OK(reader_->SetReadSchema(target_type, row_reader_options)); + pre_buffered_row_range_.reset(); return Status::OK(); } @@ -211,10 +222,20 @@ Status OrcFileBatchReader::SeekToRow(uint64_t row_number) { } Result>> OrcFileBatchReader::PreBufferRange() { - return reader_->PreBufferRange(target_column_ids_); + using ByteRanges = std::vector>; + std::pair row_range; + PAIMON_ASSIGN_OR_RAISE(ByteRanges ranges, + reader_->PreBufferRange(target_column_ids_, &row_range)); + pre_buffered_row_range_ = row_range; + return ranges; } Result OrcFileBatchReader::NextBatch() { + uint64_t next_row = reader_->GetNextRowToRead(); + if (read_ahead_ && (!pre_buffered_row_range_ || next_row < pre_buffered_row_range_->first || + next_row >= pre_buffered_row_range_->second)) { + PAIMON_RETURN_NOT_OK(read_ahead_->PrepareRanges(this)); + } PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch, reader_->Next()); if (BatchReader::IsEofBatch(batch)) { previous_batch_row_count_ = 0; diff --git a/src/paimon/format/orc/orc_file_batch_reader.h b/src/paimon/format/orc/orc_file_batch_reader.h index 43810b4d1..db7e485ea 100644 --- a/src/paimon/format/orc/orc_file_batch_reader.h +++ b/src/paimon/format/orc/orc_file_batch_reader.h @@ -21,6 +21,7 @@ #include #include #include +#include #include #include #include @@ -29,6 +30,7 @@ #include "arrow/type.h" #include "orc/OrcFile.hh" #include "orc/Reader.hh" +#include "paimon/common/reader/file_read_ahead.h" #include "paimon/format/orc/orc_format_defs.h" #include "paimon/format/orc/orc_reader_wrapper.h" #include "paimon/predicate/predicate.h" @@ -47,6 +49,11 @@ class OrcFileBatchReader : public PrefetchFileBatchReader { std::unique_ptr<::orc::InputStream>&& input_stream, const std::shared_ptr& read_memory, const std::map& options, int32_t batch_size); + static Result> Create( + std::unique_ptr<::orc::InputStream>&& input_stream, + const std::shared_ptr& read_memory, + const std::map& options, int32_t batch_size, + const std::shared_ptr& read_ahead); // For timestamp type, precision info is missing from file Result> GetFileSchema() const override; @@ -98,6 +105,8 @@ class OrcFileBatchReader : public PrefetchFileBatchReader { void Close() override { metrics_ = GetReaderMetrics(); + read_ahead_.reset(); + pre_buffered_row_range_.reset(); reader_.reset(); reader_metrics_.reset(); } @@ -111,7 +120,8 @@ class OrcFileBatchReader : public PrefetchFileBatchReader { private: OrcFileBatchReader(std::unique_ptr<::orc::ReaderMetrics>&& reader_metrics, std::unique_ptr&& reader, - const std::map& options); + const std::map& options, + const std::shared_ptr& read_ahead); static Result<::orc::RowReaderOptions> CreateRowReaderOptions( const ::orc::Type* src_type, const ::orc::Type* target_type, @@ -129,6 +139,8 @@ class OrcFileBatchReader : public PrefetchFileBatchReader { std::shared_ptr metrics_; std::vector target_column_ids_; std::vector> cache_ranges_; + std::shared_ptr read_ahead_; + std::optional> pre_buffered_row_range_; uint64_t previous_batch_row_count_ = 0; }; diff --git a/src/paimon/format/orc/orc_file_batch_reader_test.cpp b/src/paimon/format/orc/orc_file_batch_reader_test.cpp index e73e1c7ea..3cfa9b922 100644 --- a/src/paimon/format/orc/orc_file_batch_reader_test.cpp +++ b/src/paimon/format/orc/orc_file_batch_reader_test.cpp @@ -30,6 +30,8 @@ #include "arrow/c/bridge.h" #include "arrow/ipc/api.h" #include "gtest/gtest.h" +#include "paimon/common/io/cache_input_stream.h" +#include "paimon/common/metrics/metrics_impl.h" #include "paimon/common/types/data_field.h" #include "paimon/common/utils/arrow/arrow_utils.h" #include "paimon/defs.h" @@ -213,6 +215,208 @@ INSTANTIATE_TEST_SUITE_P(TestParam, OrcFileBatchReaderTest, ::testing::Values(TestParam{128 * 1024, false}, TestParam{16, false}, TestParam{16, true})); +TEST_F(OrcFileBatchReaderTest, BuilderEnablesCacheBeforeReadingMetadata) { + auto dir = paimon::test::UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + auto fs = dir->GetFileSystem(); + std::string file_path = dir->Str() + "/test.orc"; + auto schema = arrow::schema(struct_array_->type()->fields()); + WriteArray(fs, file_path, struct_array_, schema, /*options=*/{}); + ASSERT_OK_AND_ASSIGN(std::unique_ptr input, fs->Open(file_path)); + ASSERT_OK_AND_ASSIGN(int64_t file_size, input->Length()); + uint32_t factory_calls = 0; + std::shared_ptr cache; + auto stream = std::make_shared( + std::move(input), nullptr, [&]() -> Result> { + ++factory_calls; + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr cache_stream, fs->Open(file_path)); + cache = std::make_shared(cache_stream, CacheConfig(), file_size, pool_); + return cache; + }); + OrcReaderBuilder builder(/*options=*/{}, batch_size_); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, builder.Build(stream)); + ASSERT_EQ(factory_calls, 1); + ASSERT_NE(cache, nullptr); + std::shared_ptr metrics = std::make_shared(); + cache->CollectMetrics(&metrics); + ASSERT_OK_AND_ASSIGN(uint64_t block_fetches, + metrics->GetCounter(ReadAheadCacheMetrics::BLOCK_FETCHES)); + ASSERT_GT(block_fetches, 0); + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, reader->NextBatch()); + auto actual_result = arrow::ImportArray(batch.first.get(), batch.second.get()); + ASSERT_TRUE(actual_result.ok()) << actual_result.status().ToString(); + ASSERT_TRUE(actual_result.ValueOrDie()->Equals(struct_array_)); + reader->Close(); +} + +TEST_F(OrcFileBatchReaderTest, DiscoverStripeRangesOnDemand) { + using ByteRanges = std::vector>; + class CountingStream : public ::orc::InputStream { + public: + explicit CountingStream(std::unique_ptr<::orc::InputStream> inner) + : inner_(std::move(inner)) {} + uint64_t getLength() const override { + return inner_->getLength(); + } + uint64_t getNaturalReadSize() const override { + return inner_->getNaturalReadSize(); + } + const std::string& getName() const override { + return inner_->getName(); + } + void read(void* data, uint64_t size, uint64_t offset) override { + reads.emplace_back(offset, size); + inner_->read(data, size, offset); + } + std::vector> reads; + + private: + std::unique_ptr<::orc::InputStream> inner_; + }; + auto dir = paimon::test::UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + std::string path = dir->Str() + "/stripes.orc"; + { + auto output = ::orc::writeLocalFile(path); + auto type = ::orc::Type::buildTypeFromString("struct"); + auto writer = ::orc::createWriter(*type, output.get(), ::orc::WriterOptions()); + for (int64_t stripe = 0; stripe < 4; ++stripe) { + auto batch = writer->createRowBatch(4); + auto* root = dynamic_cast<::orc::StructVectorBatch*>(batch.get()); + ASSERT_NE(root, nullptr); + root->numElements = 4; + for (uint32_t column = 0; column < 2; ++column) { + auto* values = dynamic_cast<::orc::LongVectorBatch*>(root->fields[column]); + ASSERT_NE(values, nullptr); + values->numElements = 4; + for (int64_t row = 0; row < 4; ++row) { + values->data[row] = stripe * 4 + row; + } + } + writer->add(*batch); + writer->writeIntermediateFooter(); + } + writer->close(); + } + auto metadata = ::orc::createReader(::orc::readLocalFile(path), ::orc::ReaderOptions()); + ASSERT_EQ(metadata->getNumberOfStripes(), 4); + std::vector> stripe_bounds; + std::vector footer_offsets; + for (uint64_t i = 0; i < 4; ++i) { + auto stripe = metadata->getStripe(i); + uint64_t footer = stripe->getOffset() + stripe->getIndexLength() + stripe->getDataLength(); + stripe_bounds.emplace_back(stripe->getOffset(), footer); + footer_offsets.push_back(footer); + } + auto input = std::make_unique(::orc::readLocalFile(path)); + auto* observed = input.get(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr cache_stream, + dir->GetFileSystem()->Open(path)); + CacheConfig config; + config.SetBlockSize(0); + auto cache = + std::make_shared(cache_stream, config, observed->getLength(), pool_); + auto context = FileReadAhead::Create(cache, path); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + OrcFileBatchReader::Create(std::move(input), std::make_shared(pool_), {}, 2, + context)); + observed->reads.clear(); + ASSERT_OK_AND_ASSIGN(ByteRanges initial, reader->PreBufferRange()); + ASSERT_FALSE(initial.empty()); + auto assert_ranges_in_stripe = [&](const auto& ranges, uint64_t stripe) { + ASSERT_LT(stripe, stripe_bounds.size()); + ASSERT_FALSE(ranges.empty()); + for (const auto& range : ranges) { + ASSERT_GE(range.first, stripe_bounds[stripe].first); + ASSERT_LE(range.first + range.second, stripe_bounds[stripe].second); + } + }; + auto assert_no_later_footers = [&]() { + for (const auto& read : observed->reads) { + for (uint64_t i = 1; i < footer_offsets.size(); ++i) { + ASSERT_FALSE(read.first <= footer_offsets[i] && + footer_offsets[i] < read.first + read.second); + } + } + }; + assert_ranges_in_stripe(initial, 0); + assert_no_later_footers(); + ASSERT_OK(context->Warmup(reader.get())); + assert_no_later_footers(); + ASSERT_OK_AND_ASSIGN(uint64_t prepared_bytes, + context->GetMetrics()->GetCounter(ReadAheadCacheMetrics::IO_BYTES)); + ASSERT_GT(prepared_bytes, 0); + for (int64_t batch_index = 0; batch_index < 8; ++batch_index) { + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, reader->NextBatch()); + ASSERT_FALSE(BatchReader::IsEofBatch(batch)); + auto imported = arrow::ImportArray(batch.first.get(), batch.second.get()); + ASSERT_TRUE(imported.ok()); + auto root = std::dynamic_pointer_cast(imported.ValueOrDie()); + ASSERT_NE(root, nullptr); + auto values = std::dynamic_pointer_cast(root->field(0)); + ASSERT_NE(values, nullptr); + ASSERT_EQ(values->length(), 2); + ASSERT_EQ(values->Value(0), batch_index * 2); + ASSERT_OK_AND_ASSIGN(uint64_t file_row, reader->GetPreviousBatchFileRowId(0)); + ASSERT_EQ(file_row, batch_index * 2); + ASSERT_OK_AND_ASSIGN(uint64_t bytes, + context->GetMetrics()->GetCounter(ReadAheadCacheMetrics::IO_BYTES)); + if (batch_index > 0 && batch_index % 2 == 0) { + ASSERT_GT(bytes, prepared_bytes); + } else { + ASSERT_EQ(bytes, prepared_bytes); + } + prepared_bytes = bytes; + if (batch_index == 0) { + assert_no_later_footers(); + } + } + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch eof, reader->NextBatch()); + ASSERT_TRUE(BatchReader::IsEofBatch(eof)); + context->EndGeneration(); + ASSERT_OK(reader->SeekToRow(4)); + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch after_seek, reader->NextBatch()); + ASSERT_OK_AND_ASSIGN(uint64_t file_row, reader->GetPreviousBatchFileRowId(0)); + ASSERT_EQ(file_row, 4); + after_seek.first->release(after_seek.first.get()); + after_seek.second->release(after_seek.second.get()); + ASSERT_OK_AND_ASSIGN(uint64_t after_end, + context->GetMetrics()->GetCounter(ReadAheadCacheMetrics::IO_BYTES)); + ASSERT_EQ(after_end, prepared_bytes); + context->Reset(); + auto projection = arrow::schema({arrow::field("a", arrow::int32())}); + ArrowSchema projected_schema; + ASSERT_TRUE(arrow::ExportSchema(*projection, &projected_schema).ok()); + auto predicate = PredicateBuilder::GreaterOrEqual(0, "a", FieldType::INT, Literal(8)); + ASSERT_OK(reader->SetReadSchema(&projected_schema, predicate, std::nullopt)); + observed->reads.clear(); + ASSERT_OK_AND_ASSIGN(ByteRanges projected_ranges, reader->PreBufferRange()); + assert_ranges_in_stripe(projected_ranges, 0); + assert_no_later_footers(); + ASSERT_OK(context->Warmup(reader.get())); + uint64_t projected_rows = 0; + while (true) { + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, reader->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + break; + } + auto imported = arrow::ImportArray(batch.first.get(), batch.second.get()); + ASSERT_TRUE(imported.ok()); + auto root = std::dynamic_pointer_cast(imported.ValueOrDie()); + ASSERT_NE(root, nullptr); + ASSERT_EQ(root->num_fields(), 1); + auto values = std::dynamic_pointer_cast(root->field(0)); + ASSERT_NE(values, nullptr); + for (int64_t i = 0; i < values->length(); ++i) { + ASSERT_EQ(values->Value(i), 8 + projected_rows++); + } + } + ASSERT_EQ(projected_rows, 8); + reader->Close(); + context->ReleaseBuffers(); +} + TEST_F(OrcFileBatchReaderTest, TestReadBinaryWrittenFromBinaryAndLargeBinary) { auto dir = paimon::test::UniqueTestDirectory::Create(); ASSERT_TRUE(dir); diff --git a/src/paimon/format/orc/orc_reader_builder.h b/src/paimon/format/orc/orc_reader_builder.h index 603479c7d..abdfd6a1c 100644 --- a/src/paimon/format/orc/orc_reader_builder.h +++ b/src/paimon/format/orc/orc_reader_builder.h @@ -23,6 +23,7 @@ #include #include +#include "paimon/common/io/cache_input_stream.h" #include "paimon/common/utils/options_utils.h" #include "paimon/format/orc/orc_file_batch_reader.h" #include "paimon/format/orc/orc_format_defs.h" @@ -58,10 +59,16 @@ class OrcReaderBuilder : public ReaderBuilder { return Status::Invalid("natural read size should be greater than zero"); } + // Enable before parsing the footer so metadata reads also use the shared block cache. + auto cache_stream = std::dynamic_pointer_cast(path); + if (cache_stream) { + PAIMON_RETURN_NOT_OK(cache_stream->EnableCache().status()); + } PAIMON_ASSIGN_OR_RAISE(std::unique_ptr input_stream, OrcInputStreamImpl::Create(path, natural_read_size)); return OrcFileBatchReader::Create(std::move(input_stream), read_memory_, options_, - batch_size_); + batch_size_, + cache_stream ? cache_stream->GetReadAhead() : nullptr); } private: diff --git a/src/paimon/format/orc/orc_reader_wrapper.h b/src/paimon/format/orc/orc_reader_wrapper.h index cc52fde19..813b2d0ab 100644 --- a/src/paimon/format/orc/orc_reader_wrapper.h +++ b/src/paimon/format/orc/orc_reader_wrapper.h @@ -104,15 +104,40 @@ class OrcReaderWrapper { } Result>> PreBufferRange( - const std::vector& target_column_ids) { - uint64_t stripe_number = reader_->getNumberOfStripes(); - std::vector stripes; - for (uint64_t i = 0; i < stripe_number; i++) { - stripes.push_back(static_cast(i)); - } - std::list include_types(target_column_ids.begin(), target_column_ids.end()); + const std::vector& target_column_ids, std::pair* row_range) { try { - return reader_->preBufferRange(stripes, include_types); + if (stripe_row_ends_.empty()) { + uint64_t row_end = 0; + std::vector stripe_row_ends; + for (uint64_t i = 0; i < reader_->getNumberOfStripes(); ++i) { + // Stripe row counts are in the file footer; inspecting them reads no stripe + // footers. Only preBufferRange() below parses the selected stripe's footer. + uint64_t rows = reader_->getStripe(i)->getNumberOfRows(); + if (rows > std::numeric_limits::max() - row_end) { + return Status::Invalid("ORC stripe row counts overflow"); + } + row_end += rows; + stripe_row_ends.push_back(row_end); + } + stripe_row_ends_ = std::move(stripe_row_ends); + } + auto stripe = + std::upper_bound(stripe_row_ends_.begin(), stripe_row_ends_.end(), next_row_); + if (stripe == stripe_row_ends_.end()) { + if (row_range) { + *row_range = {next_row_, next_row_}; + } + return std::vector>(); + } + uint64_t index = std::distance(stripe_row_ends_.begin(), stripe); + if (index > std::numeric_limits::max()) { + return Status::Invalid("ORC stripe index exceeds uint32 bound"); + } + if (row_range) { + *row_range = {index == 0 ? 0 : stripe_row_ends_[index - 1], *stripe}; + } + std::list include_types(target_column_ids.begin(), target_column_ids.end()); + return reader_->preBufferRange({static_cast(index)}, include_types); } catch (const std::exception& e) { return Status::Invalid(fmt::format( "orc pre buffer range failed for file {}, with {} error", file_name_, e.what())); @@ -137,6 +162,7 @@ class OrcReaderWrapper { std::unique_ptr<::orc::RowReader> row_reader_; std::unique_ptr range_generator_; + std::vector stripe_row_ends_; const std::string file_name_; const int32_t batch_size_; 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 852b07027..61459f111 100644 --- a/src/paimon/format/parquet/parquet_file_batch_reader_test.cpp +++ b/src/paimon/format/parquet/parquet_file_batch_reader_test.cpp @@ -1212,6 +1212,46 @@ TEST_F(ParquetFileBatchReaderTest, TestCreateArrowReaderProperties) { } } +TEST_F(ParquetFileBatchReaderTest, BuilderUsesActualCacheBeforeChoosingPreBuffer) { + WriteArray(file_path_, struct_array_, schema_, /*write_batch_size=*/10, + /*enable_dictionary=*/false, /*max_row_group_length=*/10); + for (bool cache_enabled : {false, true}) { + SCOPED_TRACE(cache_enabled); + ParquetReaderBuilder builder({{PARQUET_READ_ENABLE_PRE_BUFFER, "true"}}, batch_size_); + ReadHints hints; + hints.read_ahead_cache_enabled = !cache_enabled; + builder.WithReadHints(hints); + ASSERT_OK_AND_ASSIGN(std::unique_ptr input, fs_->Open(file_path_)); + ASSERT_OK_AND_ASSIGN(int64_t file_size, input->Length()); + uint32_t factory_calls = 0; + CacheInputStream::CacheFactory cache_factory; + if (cache_enabled) { + cache_factory = [&]() -> Result> { + ++factory_calls; + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr cache_stream, + fs_->Open(file_path_)); + return std::make_shared(cache_stream, CacheConfig(), file_size, + GetDefaultPool()); + }; + } + auto stream = + std::make_shared(std::move(input), nullptr, std::move(cache_factory)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, builder.Build(stream)); + auto* parquet_reader = dynamic_cast(reader.get()); + ASSERT_NE(parquet_reader, nullptr); + ASSERT_EQ(parquet_reader->reader_->GetFileReader()->properties().pre_buffer(), + !cache_enabled); + ASSERT_EQ(factory_calls, cache_enabled ? 1 : 0); + ASSERT_EQ(builder.hints_->read_ahead_cache_enabled, hints.read_ahead_cache_enabled); + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, reader->NextBatch()); + auto actual_result = arrow::ImportArray(batch.first.get(), batch.second.get()); + ASSERT_TRUE(actual_result.ok()) << actual_result.status().ToString(); + std::shared_ptr actual = actual_result.ValueOrDie(); + ASSERT_TRUE(actual->Equals(struct_array_)); + reader->Close(); + } +} + TEST_F(ParquetFileBatchReaderTest, TestPreBufferReadHints) { const int32_t batch_size = 1024; // When both framework prefetch and the shared read-ahead cache are active, parquet's own diff --git a/src/paimon/format/parquet/parquet_reader_builder.h b/src/paimon/format/parquet/parquet_reader_builder.h index fad961ace..7cf792fb1 100644 --- a/src/paimon/format/parquet/parquet_reader_builder.h +++ b/src/paimon/format/parquet/parquet_reader_builder.h @@ -31,6 +31,7 @@ #include "arrow/io/memory.h" #include "fmt/format.h" #include "paimon/cache/cache.h" +#include "paimon/common/io/cache_input_stream.h" #include "paimon/common/utils/arrow/arrow_input_stream_adapter.h" #include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/format/parquet/parquet_file_batch_reader.h" @@ -81,6 +82,14 @@ class ParquetReaderBuilder : public ReaderBuilder { return Status::Invalid("Parquet reader memory pool is nullptr"); } try { + std::optional hints = hints_; + if (auto cache_stream = std::dynamic_pointer_cast(path)) { + PAIMON_ASSIGN_OR_RAISE(bool cache_active, cache_stream->EnableCache()); + if (!hints) { + hints.emplace(); + } + hints->read_ahead_cache_enabled = cache_active; + } PAIMON_ASSIGN_OR_RAISE(int64_t file_length, path->Length()); std::string file_uri; if (cache_) { @@ -92,14 +101,15 @@ class ParquetReaderBuilder : public ReaderBuilder { auto input_stream = std::make_shared(path, file_length, arrow_pool_, pool_, cache_, file_uri); auto storage_read_bytes = input_stream->StorageReadBytes(); - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<::parquet::FileMetaData> file_metadata, - GetCachedParquetMetadata(input_stream, file_uri, arrow_pool_)); + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr<::parquet::FileMetaData> file_metadata, + GetCachedParquetMetadata(input_stream, file_uri, arrow_pool_, hints)); if (file_metadata) { input_stream->SetPageIndexRanges(*file_metadata); } return ParquetFileBatchReader::Create( std::move(input_stream), options_, batch_size_, std::move(file_metadata), - std::move(storage_read_bytes), arrow_pool_, hints_); + std::move(storage_read_bytes), arrow_pool_, hints); } PAIMON_PARQUET_CATCH_AND_RETURN_STATUS("ParquetReaderBuilder::Build") } @@ -149,13 +159,14 @@ class ParquetReaderBuilder : public ReaderBuilder { Result> GetCachedParquetMetadata( const std::shared_ptr& input_stream, - const std::string& file_uri, const std::shared_ptr& arrow_pool) const { + const std::string& file_uri, const std::shared_ptr& arrow_pool, + const std::optional& hints) const { if (!cache_ || file_uri.empty()) { return std::shared_ptr<::parquet::FileMetaData>(); } PAIMON_ASSIGN_OR_RAISE( ::parquet::ReaderProperties reader_properties, - ParquetFileBatchReader::CreateReaderProperties(arrow_pool, options_, hints_)); + ParquetFileBatchReader::CreateReaderProperties(arrow_pool, options_, hints)); auto cache_key = CacheKey::ForKind(file_uri, /*position=*/-1, /*length=*/-1, CacheKind::DATA_FILE_FOOTER); diff --git a/test/inte/pk_compaction_inte_test.cpp b/test/inte/pk_compaction_inte_test.cpp index 153a7926c..2ef7635ce 100644 --- a/test/inte/pk_compaction_inte_test.cpp +++ b/test/inte/pk_compaction_inte_test.cpp @@ -140,10 +140,6 @@ class HighPositionFileBatchReader : public PrefetchFileBatchReader { return reader_->PreBufferRange(); } - void SetPreBufferRangeCallback(PreBufferRangeCallback callback) override { - reader_->SetPreBufferRangeCallback(std::move(callback)); - } - bool SupportPreciseBitmapSelection() const override { return false; } From b755d5858a2c83ec7f1e1cd131c94393d2037b3c Mon Sep 17 00:00:00 2001 From: "yonghao.fyh" Date: Thu, 5 Nov 2026 17:48:52 +0800 Subject: [PATCH 4/6] test(read): align integration fixtures with independent caching --- test/inte/read_inte_test.cpp | 15 +++++++++------ test/inte/write_and_read_inte_with_index_test.cpp | 8 ++++++-- 2 files changed, 15 insertions(+), 8 deletions(-) diff --git a/test/inte/read_inte_test.cpp b/test/inte/read_inte_test.cpp index 6f3aa0bea..8c88b4f84 100644 --- a/test/inte/read_inte_test.cpp +++ b/test/inte/read_inte_test.cpp @@ -172,8 +172,12 @@ class ReadInteTest : public testing::Test, public ::testing::WithParamInterface< EXPECT_TRUE(schema_ids.empty() || schema_ids.size() == data_files.size()); for (uint32_t j = 0; j < data_files.size(); j++) { const auto& data_file = data_files[j]; + EXPECT_OK_AND_ASSIGN( + FileStatus file_status, + LocalFileSystem().GetFileStatus(PathUtil::JoinPath(bucket_path, data_file))); data_file_metas.push_back(std::make_shared( - data_file, /*file_size=*/1, /*row_count=*/1, /*min_key=*/BinaryRow::EmptyRow(), + data_file, file_status.GetLen(), /*row_count=*/1, + /*min_key=*/BinaryRow::EmptyRow(), /*max_key=*/BinaryRow::EmptyRow(), /*key_stats=*/SimpleStats::EmptyStats(), /*value_stats=*/SimpleStats::EmptyStats(), /*min_sequence_number=*/0, /*max_sequence_number=*/0, /*schema_id=*/schema_ids.empty() ? 0 : schema_ids[j], @@ -549,11 +553,10 @@ TEST_P(ReadInteTest, TestReadAheadCacheMetrics) { ASSERT_TRUE(result_array); ASSERT_EQ(result_array->length(), 2); - // Verify the read-ahead cache metrics are surfaced through the reader chain. The prefetch - // reader merges the cache counters into its reader metrics only when a cache is created, - // so the counters must be present and effective exactly in that case. + // Cache counters must be present and effective whenever a cache is created, including + // synchronous reads with prefetch disabled. ASSERT_TRUE(read_metrics); - if (param.enable_prefetch && param.read_ahead_cache_enabled) { + if (param.read_ahead_cache_enabled) { ASSERT_OK_AND_ASSIGN(uint64_t read_count, read_metrics->GetCounter(ReadAheadCacheMetrics::READ_COUNT)); ASSERT_GT(read_count, 0u); @@ -568,7 +571,7 @@ TEST_P(ReadInteTest, TestReadAheadCacheMetrics) { ASSERT_GT(hit_bytes, 0u); ASSERT_OK(read_metrics->GetCounter(ReadAheadCacheMetrics::READ_MISSES)); ASSERT_OK(read_metrics->GetCounter(ReadAheadCacheMetrics::READ_MISS_BYTES)); - // Serving hits requires prefetch IOs issued to the underlying stream. + // Serving hits requires cache IOs issued to the underlying stream. ASSERT_OK_AND_ASSIGN(uint64_t io_count, read_metrics->GetCounter(ReadAheadCacheMetrics::IO_COUNT)); ASSERT_GT(io_count, 0u); diff --git a/test/inte/write_and_read_inte_with_index_test.cpp b/test/inte/write_and_read_inte_with_index_test.cpp index efda3de03..40aa9e889 100644 --- a/test/inte/write_and_read_inte_with_index_test.cpp +++ b/test/inte/write_and_read_inte_with_index_test.cpp @@ -1752,8 +1752,10 @@ TEST_P(WriteAndReadInteWithIndexTest, TestWithAlterTable) { auto embedded_index = std::make_shared(embedded_bytes.size(), pool_.get()); memcpy(embedded_index->data(), reinterpret_cast(embedded_bytes.data()), embedded_bytes.size()); + EXPECT_OK_AND_ASSIGN(FileStatus file_status, + LocalFileSystem().GetFileStatus(path + "bucket-0/" + file_name)); return std::make_shared( - file_name, /*file_size=*/10, + file_name, file_status.GetLen(), /*row_count=*/10, /*min_key=*/BinaryRow::EmptyRow(), /*max_key=*/BinaryRow::EmptyRow(), /*key_stats=*/SimpleStats::EmptyStats(), @@ -2431,8 +2433,10 @@ TEST_P(WriteAndReadInteWithIndexTest, TestBitmapPushDownWithMultiStripes) { } // each record is a stripe + ASSERT_OK_AND_ASSIGN(FileStatus file_status, + LocalFileSystem().GetFileStatus(path + "bucket-0/" + file_name)); auto data_file_meta = std::make_shared( - file_name, /*file_size=*/2637, + file_name, file_status.GetLen(), /*row_count=*/8, /*min_key=*/BinaryRow::EmptyRow(), /*max_key=*/BinaryRow::EmptyRow(), /*key_stats=*/SimpleStats::EmptyStats(), /*value_stats=*/SimpleStats::EmptyStats(), /*min_sequence_number=*/0, From 99621f0b443e81373278569031a9c30b2870629d Mon Sep 17 00:00:00 2001 From: "yonghao.fyh" Date: Fri, 6 Nov 2026 09:42:14 +0800 Subject: [PATCH 5/6] fix(read): preserve row mapping and prepare cache before seeks Forward synchronous batch reads through the matching underlying API so sparse bitmap compaction keeps the format reader's file-row mapping and bounds checks. Avoid opening a cache stream when LM wraps a reader without byte-range support. Prepare the target ORC stripe before seekToRow can read payload bytes. Add sparse-bitmap, LM capability and physical seek I/O regression coverage, and route on-demand ORC stripe tests through CacheInputStream. Validation: 5,430 tests passed with 2 conditional skips across seven complete test executables; changed-file pre-commit checks and git diff --check passed. --- .../reader/data_file_reader_factory_test.cpp | 49 +++++++ src/paimon/common/reader/file_read_ahead.cpp | 2 +- .../late_materializing_reader_builder.h | 5 + .../prefetch_file_batch_reader_impl.cpp | 4 +- .../reader/warmup_file_batch_reader.cpp | 53 ++++--- .../common/reader/warmup_file_batch_reader.h | 11 +- .../reader/warmup_file_batch_reader_test.cpp | 92 +++++++++++- .../format/orc/orc_file_batch_reader.cpp | 26 +++- src/paimon/format/orc/orc_file_batch_reader.h | 3 + .../format/orc/orc_file_batch_reader_test.cpp | 137 +++++++++++++++++- src/paimon/format/orc/orc_reader_wrapper.h | 7 +- 11 files changed, 339 insertions(+), 50 deletions(-) diff --git a/src/paimon/common/reader/data_file_reader_factory_test.cpp b/src/paimon/common/reader/data_file_reader_factory_test.cpp index dd5c0ffb5..f17e287b2 100644 --- a/src/paimon/common/reader/data_file_reader_factory_test.cpp +++ b/src/paimon/common/reader/data_file_reader_factory_test.cpp @@ -22,15 +22,19 @@ #include "arrow/api.h" #include "arrow/array/array_nested.h" +#include "arrow/c/bridge.h" #include "arrow/ipc/json_simple.h" #include "gtest/gtest.h" #include "paimon/common/io/cache_input_stream.h" #include "paimon/common/reader/delegating_prefetch_reader.h" +#include "paimon/common/reader/late_materializing_reader_builder.h" #include "paimon/common/reader/warmup_file_batch_reader.h" #include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/executor.h" #include "paimon/memory/memory_pool.h" +#include "paimon/predicate/literal.h" +#include "paimon/predicate/predicate_builder.h" #include "paimon/reader/file_batch_reader.h" #include "paimon/testing/mock/mock_file_system.h" #include "paimon/testing/mock/mock_format_reader_builder.h" @@ -48,6 +52,9 @@ class PlainFileBatchReader : public FileBatchReader { Result NextBatch() override { return reader_->NextBatch(); } + Result NextBatchWithBitmap() override { + return reader_->NextBatchWithBitmap(); + } std::shared_ptr GetReaderMetrics() const override { return reader_->GetReaderMetrics(); } @@ -185,6 +192,48 @@ TEST(DataFileReaderFactoryTest, CapableReaderSupportsReadAheadRegardlessOfFormat } } +TEST(DataFileReaderFactoryTest, LmPreservesPlainReaderWithoutOpeningCacheStream) { + auto type = arrow::struct_( + {arrow::field("v", arrow::int32()), arrow::field("payload", arrow::int32())}); + auto data = arrow::ipc::internal::json::ArrayFromJSON(type, "[[1,11],[2,22]]").ValueOrDie(); + auto pool = GetDefaultPool(); + for (bool cache_enabled : {false, true}) { + for (WarmupLevel level : {WarmupLevel::NONE, WarmupLevel::RAW}) { + for (bool with_predicate : {false, true}) { + auto inner = std::make_unique(data, type, false); + auto* observed = inner.get(); + LateMaterializingReaderBuilder builder(std::move(inner), GetArrowPool(pool)); + auto fs = std::make_shared(); + fs->fail_cache_open = true; + DataFileReadOptions options; + options.read_ahead_cache_enabled = cache_enabled; + options.warmup_level = level; + ASSERT_OK_AND_ASSIGN(auto reader, DataFileReaderFactory::Open( + "avro", "dummy", 0, &builder, options, fs, + nullptr, pool, GetArrowPool(pool))); + ASSERT_EQ(observed->build_count, 1); + ASSERT_EQ(fs->open_count, 1); + ASSERT_OK_AND_ASSIGN(bool cache_active, observed->stream->EnableCache()); + ASSERT_FALSE(cache_active); + ASSERT_OK_AND_ASSIGN(auto schema, reader->GetFileSchema()); + auto predicate = with_predicate ? PredicateBuilder::GreaterThan( + 0, "v", FieldType::INT, Literal(1)) + : nullptr; + ASSERT_OK(reader->SetReadSchema(schema.get(), predicate, std::nullopt)); + reader->Warmup(); + ASSERT_OK_AND_ASSIGN(auto batch, reader->NextBatch()); + ASSERT_EQ(batch.first->length, with_predicate ? 1 : 2); + auto imported = arrow::ImportArray(batch.first.get(), batch.second.get()); + ASSERT_TRUE(imported.ok()); + ASSERT_TRUE( + imported.ValueOrDie()->Equals(with_predicate ? data->Slice(1, 1) : data)); + reader->Close(); + ASSERT_EQ(fs->open_count, 1); + } + } + } +} + TEST(DataFileReaderFactoryTest, CapableReaderPropagatesCacheOpenFailureWithoutRebuilding) { auto type = arrow::struct_({arrow::field("v", arrow::int32())}); auto data = arrow::ipc::internal::json::ArrayFromJSON(type, "[[1],[2]]").ValueOrDie(); diff --git a/src/paimon/common/reader/file_read_ahead.cpp b/src/paimon/common/reader/file_read_ahead.cpp index 5322363b3..c6f5c7db9 100644 --- a/src/paimon/common/reader/file_read_ahead.cpp +++ b/src/paimon/common/reader/file_read_ahead.cpp @@ -314,7 +314,7 @@ Result> FileReadAhead::OpenSynchronous( std::shared_ptr capable_reader = checked_pointer_cast(std::move(reader)); return std::make_unique(std::move(capable_reader), std::move(state), - level, arrow_pool); + level); } Status FileReadAhead::Warmup(PrefetchFileBatchReader* reader) { diff --git a/src/paimon/common/reader/late_materializing_reader_builder.h b/src/paimon/common/reader/late_materializing_reader_builder.h index c23c220bd..917187378 100644 --- a/src/paimon/common/reader/late_materializing_reader_builder.h +++ b/src/paimon/common/reader/late_materializing_reader_builder.h @@ -57,6 +57,11 @@ class LateMaterializingReaderBuilder : public ReaderBuilder { PAIMON_ASSIGN_OR_RAISE(std::unique_ptr format_reader, inner_->Build(stream)); auto cache_stream = std::dynamic_pointer_cast(stream); + if (cache_stream && !dynamic_cast(format_reader.get())) { + // LM exposes the prefetch interface even for plain format readers. Do not let the + // outer factory enable a cache based on the wrapper's interface alone. + cache_stream->ClearCacheFactory(); + } PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader, LateMaterializingFileBatchReader::Create( std::move(format_reader), arrow_pool_, 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 93ddd1cc3..d6255c7bc 100644 --- a/src/paimon/common/reader/prefetch_file_batch_reader_impl.cpp +++ b/src/paimon/common/reader/prefetch_file_batch_reader_impl.cpp @@ -152,8 +152,8 @@ PrefetchFileBatchReaderImpl::PrefetchFileBatchReaderImpl( batch_size_(batch_size), executor_(executor), read_ahead_(read_ahead), - synchronous_reader_(std::make_unique(readers[0], read_ahead, - warmup_level, arrow_pool)), + synchronous_reader_( + std::make_unique(readers[0], read_ahead, warmup_level)), arrow_pool_(arrow_pool), prefetch_queue_capacity_(prefetch_queue_capacity), enable_adaptive_prefetch_strategy_(enable_adaptive_prefetch_strategy), diff --git a/src/paimon/common/reader/warmup_file_batch_reader.cpp b/src/paimon/common/reader/warmup_file_batch_reader.cpp index 41277180e..9c383dec5 100644 --- a/src/paimon/common/reader/warmup_file_batch_reader.cpp +++ b/src/paimon/common/reader/warmup_file_batch_reader.cpp @@ -17,6 +17,7 @@ */ #include "paimon/common/reader/warmup_file_batch_reader.h" +#include #include #include "paimon/common/metrics/metrics_impl.h" @@ -26,12 +27,8 @@ namespace paimon { WarmupFileBatchReader::WarmupFileBatchReader(std::shared_ptr reader, std::shared_ptr read_ahead, - WarmupLevel level, - std::shared_ptr arrow_pool) - : reader_(std::move(reader)), - read_ahead_(std::move(read_ahead)), - level_(level), - arrow_pool_(std::move(arrow_pool)) {} + WarmupLevel level) + : reader_(std::move(reader)), read_ahead_(std::move(read_ahead)), level_(level) {} WarmupFileBatchReader::~WarmupFileBatchReader() { ResetReadState(); @@ -55,24 +52,35 @@ void WarmupFileBatchReader::Warmup() { (void)read_ahead_->Warmup(reader_.get()); } -Result WarmupFileBatchReader::ReadOneBatch() { - PAIMON_RETURN_NOT_OK(read_ahead_->Warmup(reader_.get())); - PAIMON_ASSIGN_OR_RAISE(ReadBatchWithBitmap batch, reader_->NextBatchWithBitmap()); - Status status = read_ahead_->GetStatus(); - if (!status.ok()) { - ReaderUtils::ReleaseReadBatch(std::move(batch.first)); - return status; - } - return batch; -} - -Result WarmupFileBatchReader::NextBatchWithBitmap() { +template +Result WarmupFileBatchReader::ReadOneBatch() { if (closed_) { return Status::Invalid("cannot read a closed warmup reader"); } PAIMON_RETURN_NOT_OK(read_status_); batch_returned_ = false; - Result result = ReadOneBatch(); + Result result = [&]() -> Result { + PAIMON_RETURN_NOT_OK(read_ahead_->Warmup(reader_.get())); + // Preserve the format reader's mapping for the API used to consume this batch. + Result read = [&]() { + if constexpr (std::is_same_v) { + return reader_->NextBatch(); + } else { + return reader_->NextBatchWithBitmap(); + } + }(); + PAIMON_ASSIGN_OR_RAISE(Batch batch, std::move(read)); + Status status = read_ahead_->GetStatus(); + if (!status.ok()) { + if constexpr (std::is_same_v) { + ReaderUtils::ReleaseReadBatch(std::move(batch)); + } else { + ReaderUtils::ReleaseReadBatch(std::move(batch.first)); + } + return status; + } + return batch; + }(); if (result.ok() && !BatchReader::IsEofBatch(result.value())) { batch_returned_ = true; } else if (!result.ok()) { @@ -81,9 +89,12 @@ Result WarmupFileBatchReader::NextBatchWithBit return result; } +Result WarmupFileBatchReader::NextBatchWithBitmap() { + return ReadOneBatch(); +} + Result WarmupFileBatchReader::NextBatch() { - PAIMON_ASSIGN_OR_RAISE(ReadBatchWithBitmap batch, NextBatchWithBitmap()); - return ReaderUtils::ApplyBitmapToReadBatch(std::move(batch), arrow_pool_); + return ReadOneBatch(); } Status WarmupFileBatchReader::SetReadSchema(::ArrowSchema* schema, diff --git a/src/paimon/common/reader/warmup_file_batch_reader.h b/src/paimon/common/reader/warmup_file_batch_reader.h index 9da5645f0..bfec92127 100644 --- a/src/paimon/common/reader/warmup_file_batch_reader.h +++ b/src/paimon/common/reader/warmup_file_batch_reader.h @@ -25,18 +25,13 @@ #include "paimon/common/reader/file_read_ahead.h" #include "paimon/reader/file_batch_reader.h" -namespace arrow { -class MemoryPool; -} // namespace arrow - namespace paimon { /// A synchronous file reader with independent byte caching and raw warmup. class WarmupFileBatchReader : public FileBatchReader { public: WarmupFileBatchReader(std::shared_ptr reader, - std::shared_ptr read_ahead, WarmupLevel level, - std::shared_ptr arrow_pool); + std::shared_ptr read_ahead, WarmupLevel level); ~WarmupFileBatchReader() override; Result NextBatch() override; @@ -56,12 +51,12 @@ class WarmupFileBatchReader : public FileBatchReader { void ResetReadState(); private: - Result ReadOneBatch(); + template + Result ReadOneBatch(); std::shared_ptr reader_; std::shared_ptr read_ahead_; const WarmupLevel level_; - std::shared_ptr arrow_pool_; Status read_status_; bool prepared_ = false; bool closed_ = false; diff --git a/src/paimon/common/reader/warmup_file_batch_reader_test.cpp b/src/paimon/common/reader/warmup_file_batch_reader_test.cpp index 717b0a108..8c5e99a8a 100644 --- a/src/paimon/common/reader/warmup_file_batch_reader_test.cpp +++ b/src/paimon/common/reader/warmup_file_batch_reader_test.cpp @@ -161,8 +161,7 @@ class WarmupFileBatchReaderTest : public ::testing::Test { FileReadAhead::Open("dummy", 0, &builder, fs, 1, executor_, cache_enabled, CacheConfig(), true, pool_)); observed_ = checked_pointer_cast(opened.readers[0]); - return std::make_unique(opened.readers[0], opened.read_ahead, level, - GetArrowPool(pool_)); + return std::make_unique(opened.readers[0], opened.read_ahead, level); } Status Configure(FileBatchReader* reader) { @@ -184,6 +183,91 @@ class WarmupFileBatchReaderTest : public ::testing::Test { std::shared_ptr observed_; }; +class SparseMappingReader : public MockFileBatchReader { + public: + explicit SparseMappingReader(const std::shared_ptr& data) + : MockFileBatchReader(data, data->type(), 4) { + EnableRandomizeBatchSize(false); + } + Result NextBatchWithBitmap() override { + compact_ = false; + PAIMON_ASSIGN_OR_RAISE(ReadBatchWithBitmap batch, + MockFileBatchReader::NextBatchWithBitmap()); + if (BatchReader::IsEofBatch(batch)) { + return batch; + } + batch.second = RoaringBitmap32(); + batch.second.Add(1); + batch.second.Add(3); + return batch; + } + Result NextBatch() override { + PAIMON_ASSIGN_OR_RAISE(ReadBatchWithBitmap batch, NextBatchWithBitmap()); + PAIMON_ASSIGN_OR_RAISE( + ReadBatch result, + ReaderUtils::ApplyBitmapToReadBatch(std::move(batch), GetArrowPool(GetDefaultPool()))); + compact_ = true; + return result; + } + Result GetPreviousBatchFileRowId(uint64_t index) const override { + if (index >= (compact_ ? 2u : 4u)) { + return Status::Invalid("out of range"); + } + return compact_ ? 1 + 2 * index : index; + } + + private: + bool compact_ = false; +}; + +TEST_F(WarmupFileBatchReaderTest, PreservesSparseBitmapAndFileRowMapping) { + auto data = data_->Slice(0, 4); + for (WarmupLevel level : {WarmupLevel::NONE, WarmupLevel::RAW}) { + for (bool cache_enabled : {false, true}) { + for (bool with_bitmap : {false, true}) { + auto inner = std::make_shared(data); + std::shared_ptr cache; + if (cache_enabled) { + cache = std::make_shared(std::make_shared(), + CacheConfig(), 0, pool_); + } + WarmupFileBatchReader reader(inner, FileReadAhead::Create(cache, "sparse"), level); + ASSERT_OK(Configure(&reader)); + ASSERT_NOK(reader.GetPreviousBatchFileRowId(0)); + reader.Warmup(); + if (with_bitmap) { + ASSERT_OK_AND_ASSIGN(auto batch, reader.NextBatchWithBitmap()); + ASSERT_EQ(batch.first.first->length, 4); + ASSERT_TRUE(batch.second.Contains(1)); + ASSERT_TRUE(batch.second.Contains(3)); + ASSERT_EQ(batch.second.Cardinality(), 2); + ReaderUtils::ReleaseReadBatch(std::move(batch.first)); + } else { + ASSERT_OK_AND_ASSIGN(auto batch, reader.NextBatch()); + ASSERT_EQ(batch.first->length, 2); + auto imported = arrow::ImportArray(batch.first.get(), batch.second.get()); + ASSERT_TRUE(imported.ok()); + ASSERT_TRUE(imported.ValueOrDie()->Equals( + arrow::ipc::internal::json::ArrayFromJSON(type_, "[[1],[3]]") + .ValueOrDie())); + } + ASSERT_OK_AND_ASSIGN(uint64_t first, reader.GetPreviousBatchFileRowId(0)); + ASSERT_OK_AND_ASSIGN(uint64_t second, reader.GetPreviousBatchFileRowId(1)); + ASSERT_EQ(first, with_bitmap ? 0 : 1); + ASSERT_EQ(second, with_bitmap ? 1 : 3); + ASSERT_NOK(reader.GetPreviousBatchFileRowId(with_bitmap ? 4 : 2)); + ASSERT_OK_AND_ASSIGN(auto eof, reader.NextBatch()); + ASSERT_TRUE(BatchReader::IsEofBatch(eof)); + ASSERT_NOK(reader.GetPreviousBatchFileRowId(0)); + ASSERT_OK(Configure(&reader)); + ASSERT_NOK(reader.GetPreviousBatchFileRowId(0)); + reader.Close(); + ASSERT_NOK(reader.GetPreviousBatchFileRowId(0)); + } + } + } +} + TEST_F(WarmupFileBatchReaderTest, RawWarmsBytesWithoutDecodingAndIsIdempotent) { auto fs = std::make_shared(); ASSERT_OK_AND_ASSIGN(auto reader, Open(WarmupLevel::RAW, true, fs)); @@ -378,8 +462,8 @@ TEST_F(WarmupFileBatchReaderTest, InvalidInitialAndLateRangesReportFileOnRead) { FileReadAhead::Open("range-bounds-file", late ? 4 : 2, &builder, std::make_shared(), 1, nullptr, true, CacheConfig(), false, pool_)); - auto reader = std::make_unique( - opened.readers[0], opened.read_ahead, WarmupLevel::RAW, GetArrowPool(pool_)); + auto reader = std::make_unique(opened.readers[0], opened.read_ahead, + WarmupLevel::RAW); ASSERT_OK(Configure(reader.get())); reader->Warmup(); if (late) { diff --git a/src/paimon/format/orc/orc_file_batch_reader.cpp b/src/paimon/format/orc/orc_file_batch_reader.cpp index 91fb931ab..937cb7577 100644 --- a/src/paimon/format/orc/orc_file_batch_reader.cpp +++ b/src/paimon/format/orc/orc_file_batch_reader.cpp @@ -218,24 +218,36 @@ Status OrcFileBatchReader::SetReadSchema(::ArrowSchema* read_schema, } Status OrcFileBatchReader::SeekToRow(uint64_t row_number) { + // ORC may read payload while seeking within a stripe. Warm the target before that I/O. + PAIMON_RETURN_NOT_OK(PrepareRangesForRow(row_number)); return reader_->SeekToRow(row_number); } Result>> OrcFileBatchReader::PreBufferRange() { using ByteRanges = std::vector>; std::pair row_range; - PAIMON_ASSIGN_OR_RAISE(ByteRanges ranges, - reader_->PreBufferRange(target_column_ids_, &row_range)); + PAIMON_ASSIGN_OR_RAISE( + ByteRanges ranges, + reader_->PreBufferRange(target_column_ids_, + range_preparation_row_.value_or(reader_->GetNextRowToRead()), + &row_range)); pre_buffered_row_range_ = row_range; return ranges; } -Result OrcFileBatchReader::NextBatch() { - uint64_t next_row = reader_->GetNextRowToRead(); - if (read_ahead_ && (!pre_buffered_row_range_ || next_row < pre_buffered_row_range_->first || - next_row >= pre_buffered_row_range_->second)) { - PAIMON_RETURN_NOT_OK(read_ahead_->PrepareRanges(this)); +Status OrcFileBatchReader::PrepareRangesForRow(uint64_t row_number) { + if (read_ahead_ && (!pre_buffered_row_range_ || row_number < pre_buffered_row_range_->first || + row_number >= pre_buffered_row_range_->second)) { + range_preparation_row_ = row_number; + Status status = read_ahead_->PrepareRanges(this); + range_preparation_row_.reset(); + return status; } + return Status::OK(); +} + +Result OrcFileBatchReader::NextBatch() { + PAIMON_RETURN_NOT_OK(PrepareRangesForRow(reader_->GetNextRowToRead())); PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch, reader_->Next()); if (BatchReader::IsEofBatch(batch)) { previous_batch_row_count_ = 0; diff --git a/src/paimon/format/orc/orc_file_batch_reader.h b/src/paimon/format/orc/orc_file_batch_reader.h index db7e485ea..baa9f0bc0 100644 --- a/src/paimon/format/orc/orc_file_batch_reader.h +++ b/src/paimon/format/orc/orc_file_batch_reader.h @@ -118,6 +118,8 @@ class OrcFileBatchReader : public PrefetchFileBatchReader { Result>> PreBufferRange() override; private: + Status PrepareRangesForRow(uint64_t row_number); + OrcFileBatchReader(std::unique_ptr<::orc::ReaderMetrics>&& reader_metrics, std::unique_ptr&& reader, const std::map& options, @@ -141,6 +143,7 @@ class OrcFileBatchReader : public PrefetchFileBatchReader { std::vector> cache_ranges_; std::shared_ptr read_ahead_; std::optional> pre_buffered_row_range_; + std::optional range_preparation_row_; uint64_t previous_batch_row_count_ = 0; }; diff --git a/src/paimon/format/orc/orc_file_batch_reader_test.cpp b/src/paimon/format/orc/orc_file_batch_reader_test.cpp index 3cfa9b922..e086d596d 100644 --- a/src/paimon/format/orc/orc_file_batch_reader_test.cpp +++ b/src/paimon/format/orc/orc_file_batch_reader_test.cpp @@ -249,6 +249,130 @@ TEST_F(OrcFileBatchReaderTest, BuilderEnablesCacheBeforeReadingMetadata) { reader->Close(); } +class PhysicalProbeStream : public paimon::InputStream { + public: + explicit PhysicalProbeStream(std::shared_ptr inner) + : inner_(std::move(inner)) {} + Status Seek(int64_t offset, SeekOrigin origin) override { + return inner_->Seek(offset, origin); + } + Result GetPos() const override { + return inner_->GetPos(); + } + Result Length() const override { + return inner_->Length(); + } + Result GetUri() const override { + return inner_->GetUri(); + } + Status Close() override { + return inner_->Close(); + } + Result Read(char* data, int64_t size) override { + return inner_->Read(data, size); + } + Result Read(char* data, int64_t size, int64_t offset) override { + reads.emplace_back(offset, size); + return inner_->Read(data, size, offset); + } + void ReadAsync(char* data, int64_t size, int64_t offset, + std::function&& cb) override { + reads.emplace_back(offset, size); + inner_->ReadAsync(data, size, offset, std::move(cb)); + } + uint64_t Overlap(uint64_t begin, uint64_t length) const { + uint64_t bytes = 0; + for (const auto& r : reads) { + uint64_t lo = std::max(begin, r.first); + uint64_t hi = std::min(begin + length, r.first + r.second); + if (hi > lo) { + bytes += hi - lo; + } + } + return bytes; + } + std::vector> reads; + + private: + std::shared_ptr inner_; +}; + +TEST_F(OrcFileBatchReaderTest, SeekPreparesTargetStripeBeforePhysicalReads) { + auto dir = paimon::test::UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + auto path = dir->Str() + "/seek.orc"; + { + auto output = ::orc::writeLocalFile(path); + auto type = ::orc::Type::buildTypeFromString("struct"); + auto writer = ::orc::createWriter(*type, output.get(), ::orc::WriterOptions()); + for (int stripe = 0; stripe < 3; ++stripe) { + auto batch = writer->createRowBatch(4); + auto* root = dynamic_cast<::orc::StructVectorBatch*>(batch.get()); + ASSERT_NE(root, nullptr); + root->numElements = 4; + auto* values = dynamic_cast<::orc::LongVectorBatch*>(root->fields[0]); + ASSERT_NE(values, nullptr); + values->numElements = 4; + for (int row = 0; row < 4; ++row) { + values->data[row] = stripe * 4 + row; + } + writer->add(*batch); + writer->writeIntermediateFooter(); + } + writer->close(); + } + auto metadata = ::orc::createReader(::orc::readLocalFile(path), ::orc::ReaderOptions()); + ASSERT_EQ(metadata->getNumberOfStripes(), 3); + auto stripe = metadata->getStripe(2); + uint64_t payload_offset = stripe->getOffset() + stripe->getIndexLength(); + uint64_t payload_length = stripe->getDataLength(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr raw, + dir->GetFileSystem()->Open(path)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr cache_raw, + dir->GetFileSystem()->Open(path)); + auto physical_owner = std::make_unique(raw); + auto* physical = physical_owner.get(); + auto cache_physical = std::make_shared(cache_raw); + auto pool = GetDefaultPool(); + CacheConfig config; + config.SetBlockSize(0); + ASSERT_OK_AND_ASSIGN(int64_t file_size, raw->Length()); + auto cache = std::make_shared(cache_physical, config, file_size, pool); + auto context = FileReadAhead::Create(cache, path); + auto cache_input = std::make_shared( + std::move(physical_owner), cache, CacheInputStream::CacheFactory(), context); + ASSERT_OK_AND_ASSIGN(auto orc_input, OrcInputStreamImpl::Create(cache_input, 1)); + ASSERT_OK_AND_ASSIGN(auto reader, OrcFileBatchReader::Create( + std::move(orc_input), + std::make_shared(pool), {}, 2, context)); + ASSERT_OK(context->Warmup(reader.get())); + physical->reads.clear(); + cache_physical->reads.clear(); + ASSERT_OK_AND_ASSIGN(uint64_t hits_before, + context->GetMetrics()->GetCounter(ReadAheadCacheMetrics::READ_HITS)); + ASSERT_OK(reader->SeekToRow(9)); + uint64_t seek_bytes = physical->Overlap(payload_offset, payload_length); + uint64_t seek_cache_bytes = cache_physical->Overlap(payload_offset, payload_length); + ASSERT_OK_AND_ASSIGN(auto batch, reader->NextBatch()); + uint64_t next_cache_bytes = cache_physical->Overlap(payload_offset, payload_length); + ASSERT_EQ(seek_bytes, 0); + ASSERT_EQ(seek_cache_bytes, payload_length); + ASSERT_EQ(next_cache_bytes, seek_cache_bytes); + ASSERT_OK_AND_ASSIGN(uint64_t hits, + context->GetMetrics()->GetCounter(ReadAheadCacheMetrics::READ_HITS)); + ASSERT_GT(hits, hits_before); + ASSERT_EQ(physical->Overlap(payload_offset, payload_length), 0); + ASSERT_OK_AND_ASSIGN(uint64_t file_row, reader->GetPreviousBatchFileRowId(0)); + ASSERT_EQ(file_row, 9); + auto imported = arrow::ImportArray(batch.first.get(), batch.second.get()); + ASSERT_TRUE(imported.ok()); + auto values = std::static_pointer_cast( + std::static_pointer_cast(imported.ValueOrDie())->field(0)); + ASSERT_EQ(values->Value(0), 9); + reader->Close(); + context->ReleaseBuffers(); +} + TEST_F(OrcFileBatchReaderTest, DiscoverStripeRangesOnDemand) { using ByteRanges = std::vector>; class CountingStream : public ::orc::InputStream { @@ -308,15 +432,20 @@ TEST_F(OrcFileBatchReaderTest, DiscoverStripeRangesOnDemand) { stripe_bounds.emplace_back(stripe->getOffset(), footer); footer_offsets.push_back(footer); } - auto input = std::make_unique(::orc::readLocalFile(path)); - auto* observed = input.get(); + ASSERT_OK_AND_ASSIGN(std::unique_ptr main_stream, + dir->GetFileSystem()->Open(path)); + ASSERT_OK_AND_ASSIGN(int64_t file_size, main_stream->Length()); ASSERT_OK_AND_ASSIGN(std::shared_ptr cache_stream, dir->GetFileSystem()->Open(path)); CacheConfig config; config.SetBlockSize(0); - auto cache = - std::make_shared(cache_stream, config, observed->getLength(), pool_); + auto cache = std::make_shared(cache_stream, config, file_size, pool_); auto context = FileReadAhead::Create(cache, path); + auto cached_stream = std::make_shared( + std::move(main_stream), cache, CacheInputStream::CacheFactory(), context); + ASSERT_OK_AND_ASSIGN(auto adapter, OrcInputStreamImpl::Create(cached_stream, 1)); + auto input = std::make_unique(std::move(adapter)); + auto* observed = input.get(); ASSERT_OK_AND_ASSIGN( std::unique_ptr reader, OrcFileBatchReader::Create(std::move(input), std::make_shared(pool_), {}, 2, diff --git a/src/paimon/format/orc/orc_reader_wrapper.h b/src/paimon/format/orc/orc_reader_wrapper.h index 813b2d0ab..c7682dd4e 100644 --- a/src/paimon/format/orc/orc_reader_wrapper.h +++ b/src/paimon/format/orc/orc_reader_wrapper.h @@ -104,7 +104,8 @@ class OrcReaderWrapper { } Result>> PreBufferRange( - const std::vector& target_column_ids, std::pair* row_range) { + const std::vector& target_column_ids, uint64_t next_row, + std::pair* row_range) { try { if (stripe_row_ends_.empty()) { uint64_t row_end = 0; @@ -122,10 +123,10 @@ class OrcReaderWrapper { stripe_row_ends_ = std::move(stripe_row_ends); } auto stripe = - std::upper_bound(stripe_row_ends_.begin(), stripe_row_ends_.end(), next_row_); + std::upper_bound(stripe_row_ends_.begin(), stripe_row_ends_.end(), next_row); if (stripe == stripe_row_ends_.end()) { if (row_range) { - *row_range = {next_row_, next_row_}; + *row_range = {next_row, next_row}; } return std::vector>(); } From 7e32d5f8d896de8a23ec401bc06e9086e757f3a7 Mon Sep 17 00:00:00 2001 From: "yonghao.fyh" Date: Fri, 6 Nov 2026 10:54:40 +0800 Subject: [PATCH 6/6] test(read): use filesystem factory for fixture metadata --- test/inte/write_and_read_inte_with_index_test.cpp | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/test/inte/write_and_read_inte_with_index_test.cpp b/test/inte/write_and_read_inte_with_index_test.cpp index 40aa9e889..209a12f19 100644 --- a/test/inte/write_and_read_inte_with_index_test.cpp +++ b/test/inte/write_and_read_inte_with_index_test.cpp @@ -56,7 +56,7 @@ #include "paimon/data/timestamp.h" #include "paimon/defs.h" #include "paimon/factories/factory_creator.h" -#include "paimon/fs/local/local_file_system.h" +#include "paimon/fs/file_system_factory.h" #include "paimon/memory/bytes.h" #include "paimon/memory/memory_pool.h" #include "paimon/metrics.h" @@ -1747,13 +1747,15 @@ TEST_P(WriteAndReadInteWithIndexTest, TestWithAlterTable) { file_name2 = "data-2efd8477-9fb4-44a6-9a39-901dd5a83b28-0.parquet"; } + ASSERT_OK_AND_ASSIGN(std::unique_ptr file_system, + FileSystemFactory::Get("local", path, {})); auto create_data_file_meta = [&](const std::string& file_name, int32_t schema_id, const std::vector& embedded_bytes) { auto embedded_index = std::make_shared(embedded_bytes.size(), pool_.get()); memcpy(embedded_index->data(), reinterpret_cast(embedded_bytes.data()), embedded_bytes.size()); EXPECT_OK_AND_ASSIGN(FileStatus file_status, - LocalFileSystem().GetFileStatus(path + "bucket-0/" + file_name)); + file_system->GetFileStatus(path + "bucket-0/" + file_name)); return std::make_shared( file_name, file_status.GetLen(), /*row_count=*/10, /*min_key=*/BinaryRow::EmptyRow(), @@ -2433,8 +2435,10 @@ TEST_P(WriteAndReadInteWithIndexTest, TestBitmapPushDownWithMultiStripes) { } // each record is a stripe + ASSERT_OK_AND_ASSIGN(std::unique_ptr file_system, + FileSystemFactory::Get("local", path, {})); ASSERT_OK_AND_ASSIGN(FileStatus file_status, - LocalFileSystem().GetFileStatus(path + "bucket-0/" + file_name)); + file_system->GetFileStatus(path + "bucket-0/" + file_name)); auto data_file_meta = std::make_shared( file_name, file_status.GetLen(), /*row_count=*/8, /*min_key=*/BinaryRow::EmptyRow(),