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..1eb96b189 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,25 +91,29 @@ 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 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 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 +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. 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 @@ -118,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/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..cb06e05a7 100644 --- a/include/paimon/read_context.h +++ b/include/paimon/read_context.h @@ -414,26 +414,27 @@ 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. + /// 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); - /// 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. 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. + /// 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/reader/prefetch_file_batch_reader.h b/include/paimon/reader/prefetch_file_batch_reader.h index 1969194bf..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 @@ -48,7 +47,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 { @@ -136,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/include/paimon/utils/prefetch_cache_config.h b/include/paimon/utils/prefetch_cache_config.h index 10a82f1f8..e06cc129f 100644 --- a/include/paimon/utils/prefetch_cache_config.h +++ b/include/paimon/utils/prefetch_cache_config.h @@ -124,26 +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. - DECODED, }; } // namespace paimon 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/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/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 edaac8301..2bba5fff8 100644 --- a/src/paimon/common/reader/data_file_reader_factory.cpp +++ b/src/paimon/common/reader/data_file_reader_factory.cpp @@ -22,6 +22,7 @@ #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/format/file_format.h" #include "paimon/format/file_format_factory.h" @@ -69,7 +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_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; } @@ -92,9 +95,10 @@ Result> DataFileReaderFactory::Open( read_options.warmup_level, pool, arrow_pool)); return std::make_unique(std::move(prefetch_reader)); } - 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.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..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,12 +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" @@ -36,6 +43,211 @@ namespace paimon::test { +// 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(); + } + Result NextBatchWithBitmap() override { + return reader_->NextBatchWithBitmap(); + } + 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(); + auto pool = GetDefaultPool(); + DataFileReadOptions options; + 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_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, 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(); + 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. @@ -53,7 +265,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/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..c6f5c7db9 --- /dev/null +++ b/src/paimon/common/reader/file_read_ahead.cpp @@ -0,0 +1,446 @@ +/* + * 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 "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 { + +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 + +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, + 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); + } + 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; + 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, + 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, CacheInputStream::CacheFactory(), state); + 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); + } + 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); +} + +Status FileReadAhead::Warmup(PrefetchFileBatchReader* reader) { + std::lock_guard lock(mutex_); + if (!cache_ || warmed_) { + return status_; + } + warmed_ = true; + 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_; + 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_ = Status(added.status().code(), fmt::format("read-ahead ranges for file '{}': {}", + path_, added.status().ToString())); + ++error_count_; + return status_; + } + if (initial) { + cache_->Warmup(); + } else if (added.value()) { + cache_->Warmup(added.value().value()); + } + return Status::OK(); +} + +void FileReadAhead::SetSynchronous(bool synchronous) { + if (cache_) { + cache_->SetBoundedRetention(synchronous); + } +} + +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..4a4dfa4b5 --- /dev/null +++ b/src/paimon/common/reader/file_read_ahead.h @@ -0,0 +1,96 @@ +/* + * 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" +#include "paimon/visibility.h" + +namespace arrow { +class MemoryPool; +} // namespace arrow + +namespace paimon { + +class Executor; +class ReaderBuilder; +struct PrefetchIoMetricsState; + +/// 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 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, + const std::shared_ptr& executor, bool cache_enabled, + 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(); + void Reset(); + void ReleaseBuffers(); + std::shared_ptr GetMetrics() const; + + private: + FileReadAhead(const std::shared_ptr& cache, const std::string& path) + : cache_(cache), path_(path) {} + + 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_; + uint64_t error_count_ = 0; +}; + +} // namespace paimon 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 a95fd91a9..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" @@ -706,7 +709,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 +745,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 +808,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; @@ -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}); @@ -960,7 +1140,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/late_materializing_reader_builder.h b/src/paimon/common/reader/late_materializing_reader_builder.h index 56571fcb9..917187378 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,16 @@ 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); + 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_, + 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 4e5e540e4..d6255c7bc 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), 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)); @@ -366,16 +169,13 @@ PrefetchFileBatchReaderImpl::PrefetchFileBatchReaderImpl( PrefetchFileBatchReaderImpl::~PrefetchFileBatchReaderImpl() { (void)CleanUp(); - ClearPreBufferRangeCallbacks(); } 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 +190,7 @@ Status PrefetchFileBatchReaderImpl::SetReadSchema( Status PrefetchFileBatchReaderImpl::RefreshReadRanges() { PAIMON_RETURN_NOT_OK(CleanUp()); - if (cache_) { - cache_->Reset(); - } + read_ahead_->Reset(); return RefreshReadRangesAfterCleanUp(); } @@ -415,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 = @@ -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_->ResetReadState(); 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++) { @@ -539,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) { @@ -776,74 +572,13 @@ 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()); -} - -void PrefetchFileBatchReaderImpl::ClearPreBufferRangeCallbacks() { - // Belt and braces: the callbacks capture this, and by now CleanUp() has joined the background - // thread, so no sub-reader can report a range any more. - for (const auto& sub_reader : readers_) { - sub_reader->SetPreBufferRangeCallback(nullptr); - } + // 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::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, @@ -853,20 +588,16 @@ void PrefetchFileBatchReaderImpl::Warmup() { if (!read_ranges_freshed_) { 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() { 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 +684,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 +697,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 +715,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 +728,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 " @@ -1045,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_); @@ -1054,7 +764,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,16 +783,10 @@ 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(); } } // 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 164cf6c6b..c7de0c126 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,37 +127,16 @@ 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. + /// Starts the background batch prefetch thread lazily on the first read. 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(); void SetReadStatus(const Status& status); Status GetReadStatus() const; Result IsEofRange(const std::pair& read_range) const; @@ -196,7 +179,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 +190,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..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 @@ -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" @@ -359,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; } @@ -395,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)); @@ -423,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, @@ -477,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"); @@ -495,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(); @@ -742,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}}; @@ -773,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(); @@ -786,67 +783,77 @@ 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)); } -// 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()); - - DelegatingPrefetchReader delegating_reader(std::move(reader)); - delegating_reader.Warmup(); - 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)); +// 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); + 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::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_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_FALSE(DecodeThreadStarted(prefetch_reader)); + delegating_reader.Close(); + delegating_reader.Warmup(); + ASSERT_NOK(delegating_reader.NextBatchWithBitmap()); + } } TEST_F(PrefetchFileBatchReaderImplTest, SetReadRanges) { @@ -861,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; @@ -905,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; @@ -931,14 +938,17 @@ 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(); 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) { @@ -953,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; @@ -972,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_ = {{}}; @@ -991,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; @@ -1012,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()) @@ -1058,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++) { @@ -1097,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; @@ -1117,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; @@ -1161,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) { @@ -1176,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( @@ -1184,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( @@ -1193,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"); } { @@ -1203,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( @@ -1211,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( @@ -1219,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( @@ -1229,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"); } @@ -1255,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; @@ -1301,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())); @@ -1333,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))); @@ -1365,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())); @@ -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}) { + 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}) { + 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); @@ -1449,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); @@ -1526,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 @@ -1583,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 new file mode 100644 index 000000000..9c383dec5 --- /dev/null +++ b/src/paimon/common/reader/warmup_file_batch_reader.cpp @@ -0,0 +1,152 @@ +/* + * 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 "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) + : reader_(std::move(reader)), read_ahead_(std::move(read_ahead)), level_(level) {} + +WarmupFileBatchReader::~WarmupFileBatchReader() { + ResetReadState(); + read_ahead_->ReleaseBuffers(); +} + +void WarmupFileBatchReader::ResetReadState() { + prepared_ = false; + batch_returned_ = false; + read_status_ = Status::OK(); +} + +void WarmupFileBatchReader::Prepare() { + prepared_ = true; +} + +void WarmupFileBatchReader::Warmup() { + if (!prepared_ || closed_ || level_ != WarmupLevel::RAW) { + return; + } + (void)read_ahead_->Warmup(reader_.get()); +} + +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 = [&]() -> 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()) { + read_status_ = result.status(); + } + return result; +} + +Result WarmupFileBatchReader::NextBatchWithBitmap() { + return ReadOneBatch(); +} + +Result WarmupFileBatchReader::NextBatch() { + return ReadOneBatch(); +} + +Status WarmupFileBatchReader::SetReadSchema(::ArrowSchema* schema, + const std::shared_ptr& predicate, + const std::optional& bitmap) { + ResetReadState(); + 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 { + auto metrics = std::make_shared(); + 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 { + 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 { + return reader_->GetNumberOfRows(); +} + +bool WarmupFileBatchReader::SupportPreciseBitmapSelection() const { + return reader_->SupportPreciseBitmapSelection(); +} + +void WarmupFileBatchReader::Close() { + if (closed_) { + return; + } + ResetReadState(); + closed_ = true; + 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..bfec92127 --- /dev/null +++ b/src/paimon/common/reader/warmup_file_batch_reader.h @@ -0,0 +1,66 @@ +/* + * 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 "arrow/c/abi.h" +#include "paimon/common/reader/file_read_ahead.h" +#include "paimon/reader/file_batch_reader.h" + +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); + ~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 ResetReadState(); + + private: + template + Result ReadOneBatch(); + + std::shared_ptr reader_; + std::shared_ptr read_ahead_; + const WarmupLevel level_; + Status read_status_; + bool prepared_ = 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..8c5e99a8a --- /dev/null +++ b/src/paimon/common/reader/warmup_file_batch_reader_test.cpp @@ -0,0 +1,499 @@ +/* + * 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 "arrow/api.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/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" +#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); + if (auto cache_stream = std::dynamic_pointer_cast(stream)) { + read_ahead_ = cache_stream->GetReadAhead(); + } + } + + Result NextBatchWithBitmap() override { + ++read_count; + last_read_thread = std::this_thread::get_id(); + return Read(); + } + + Result>> PreBufferRange() override { + PAIMON_RETURN_NOT_OK(range_status); + return std::vector>{{range_offset_, range_length}}; + } + + 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, (read_count - 1) * read_stride)); + if (size != 4) { + return Status::IOError("short observed read"); + } + } + return MockFileBatchReader::NextBatchWithBitmap(); + } + + std::shared_ptr stream_; + std::shared_ptr read_ahead_; + uint64_t range_offset_ = 0; +}; + +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 { + auto reader = std::make_unique(data_, stream); + last_reader = reader.get(); + return reader; + } + + mutable ObservedReader* last_reader = nullptr; + + 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); + } + + 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_; +}; + +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)); + 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, 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, 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, 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, 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_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, 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(); + 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, EmptyFileAndRawWarmupCloseSafely) { + for (bool empty : {false, true}) { + if (empty) { + data_ = data_->Slice(0, 0); + } + ASSERT_OK_AND_ASSIGN(auto reader, Open(WarmupLevel::RAW, 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::RAW, 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, 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, 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); + 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())); + ASSERT_OK(observed_->PrepareLateRanges(16)); + ASSERT_EQ(Counter(reader.get(), ReadAheadCacheMetrics::IO_COUNT), 0); + reader->Warmup(); + ASSERT_OK(observed_->PrepareLateRanges(16)); + ASSERT_EQ(Counter(reader.get(), ReadAheadCacheMetrics::IO_COUNT), 2); + ASSERT_OK(observed_->PrepareLateRanges(16)); + ASSERT_EQ(Counter(reader.get(), ReadAheadCacheMetrics::IO_COUNT), 2); + ASSERT_OK(Configure(reader.get())); + 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(); + ASSERT_OK(observed_->PrepareLateRanges(32)); + ASSERT_EQ(Counter(reader.get(), ReadAheadCacheMetrics::IO_COUNT), 1); +} + +} // namespace paimon::test 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/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/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/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/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/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()); diff --git a/src/paimon/format/orc/orc_file_batch_reader.cpp b/src/paimon/format/orc/orc_file_batch_reader.cpp index 526ff7cd2..937cb7577 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,18 +213,41 @@ 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(); } 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() { - 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_, + range_preparation_row_.value_or(reader_->GetNextRowToRead()), + &row_range)); + pre_buffered_row_range_ = row_range; + return ranges; +} + +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 43810b4d1..baa9f0bc0 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(); } @@ -109,9 +118,12 @@ 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); + 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 +141,9 @@ 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_; + 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 e73e1c7ea..e086d596d 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,337 @@ 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(); +} + +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 { + 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); + } + 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, 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, + 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..c7682dd4e 100644 --- a/src/paimon/format/orc/orc_reader_wrapper.h +++ b/src/paimon/format/orc/orc_reader_wrapper.h @@ -104,15 +104,41 @@ 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, uint64_t next_row, + 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 +163,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.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..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 @@ -1239,7 +1279,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 +1287,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. { 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; } 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..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,17 @@ 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, + file_system->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 +2435,12 @@ 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, + file_system->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,