From 907f440904a0b0f1ccfc0d63e62fe0dc682912ad Mon Sep 17 00:00:00 2001 From: Justin Li Date: Mon, 13 Apr 2026 14:06:10 -0400 Subject: [PATCH 01/12] GH-39808: [C++][Parquet] Evict pre-buffered row-group bytes after decode Dataset.to_batches() on parquet files accumulates memory as iteration proceeds because ReadRangeCache has no eviction API. PreBuffer() is called once with every row group up front, entries stay resident until the FileReader is destroyed, and users see roughly 10x more memory than the equivalent ParquetFile.iter_batches() path. This is one of the longest-standing open issues on the tracker. Add a new ReadRangeCache::EvictEntriesInRange(start, length) method that removes cache entries fully contained in the given window. Entries that span past the window (for example, because range coalescing merged them with an adjacent row group's column chunk) are deliberately left in place, so eviction is safe in the presence of coalescing. Expose the primitive through ParquetFileReader::EvictPreBufferedData and call it from the Arrow RowGroupGenerator's .Then callback once a row group has been decoded into Arrow arrays. At that point the raw column-chunk bytes held by the cache are no longer needed, and releasing them gives each row group a bounded per-row-group memory footprint. Thread safety: promote the existing mutex from LazyImpl into base Impl so that Cache, Read, Wait, WaitFor, and EvictEntriesInRange all acquire it before touching the entries vector. Concurrent Read from one thread and Evict from another was previously undefined behaviour in the non-lazy cache, and the dataset scanner's batch_readahead path is exactly the concurrent call pattern that would trigger it. Read now drops the lock before blocking on the I/O future, so the new locking does not serialize readers more tightly than before. Measured on a 458 MB / 10-row-group / 10M-row test file: Dataset.to_batches, before fix: 598 MB peak Dataset.to_batches, after fix: 331 MB peak (-267 MB) ParquetFile.iter_batches (no pre-buffer): 59 MB peak Savings scale linearly with row-group count, so on the multi-GB files from the issue thread this single fix recovers several GB of peak allocation. The remaining gap between Dataset.to_batches and iter_batches comes from a second source of accumulation in the Dataset infrastructure that is unrelated to the ReadRangeCache and should be tracked as a follow-up issue. New tests: * RangeReadCache.EvictEntriesInRange * RangeReadCache.EvictEntriesInRangeSpanningEntry * RangeReadCache.ConcurrentReadAndEvict * TestArrowReadWrite.EvictPreBufferedData * TestArrowReadWrite.GetRecordBatchGeneratorReleasesPreBufferedRowGroups Full regression sweep: 824/824 parquet-arrow-reader-writer-test, 57/57 arrow-io-memory-test. --- cpp/src/arrow/io/caching.cc | 183 +++++++++++------- cpp/src/arrow/io/caching.h | 15 ++ cpp/src/arrow/io/memory_test.cc | 182 +++++++++++++++++ .../parquet/arrow/arrow_reader_writer_test.cc | 110 +++++++++++ cpp/src/parquet/arrow/reader.cc | 12 +- cpp/src/parquet/file_reader.cc | 44 +++++ cpp/src/parquet/file_reader.h | 17 ++ 7 files changed, 494 insertions(+), 69 deletions(-) diff --git a/cpp/src/arrow/io/caching.cc b/cpp/src/arrow/io/caching.cc index 74e98170ad0b..abb1cafcb059 100644 --- a/cpp/src/arrow/io/caching.cc +++ b/cpp/src/arrow/io/caching.cc @@ -151,11 +151,19 @@ struct ReadRangeCache::Impl { IOContext ctx; CacheOptions options; - // Ordered by offset (so as to find a matching region by binary search) + // Ordered by offset (so as to find a matching region by binary search). + // Mutation of `entries` and of individual entries' futures must be + // serialized via `entry_mutex`. Every public method that touches either + // acquires the mutex before delegating to the protected *Locked helpers + // below, so both the eager and lazy variants are safe to call concurrently + // from multiple threads. std::vector entries; + std::mutex entry_mutex; virtual ~Impl() = default; + // -- Polymorphic hooks. Always called with entry_mutex held. -- + // Get the future corresponding to a range virtual Future> MaybeRead(RangeCacheEntry* entry) { return entry->future; @@ -172,23 +180,31 @@ struct ReadRangeCache::Impl { return new_entries; } - // Add the given ranges to the cache, coalescing them where possible - virtual Status Cache(std::vector ranges) { + // -- Public entry points (acquire entry_mutex, then delegate). -- + + // Add the given ranges to the cache, coalescing them where possible. + Status Cache(std::vector ranges) { ARROW_ASSIGN_OR_RAISE( ranges, internal::CoalesceReadRanges(std::move(ranges), options.hole_size_limit, options.range_size_limit)); - std::vector new_entries = MakeCacheEntries(ranges); - // Add new entries, themselves ordered by offset - if (entries.size() > 0) { - std::vector merged(entries.size() + new_entries.size()); - std::merge(entries.begin(), entries.end(), new_entries.begin(), new_entries.end(), - merged.begin()); - entries = std::move(merged); - } else { - entries = std::move(new_entries); + Status st; + { + std::unique_lock guard(entry_mutex); + std::vector new_entries = MakeCacheEntries(ranges); + // Add new entries, themselves ordered by offset + if (entries.size() > 0) { + std::vector merged(entries.size() + new_entries.size()); + std::merge(entries.begin(), entries.end(), new_entries.begin(), + new_entries.end(), merged.begin()); + entries = std::move(merged); + } else { + entries = std::move(new_entries); + } } - // Prefetch immediately, regardless of executor availability, if possible - auto st = file->WillNeed(ranges); + // Prefetch immediately, regardless of executor availability, if possible. + // Do this outside the lock: WillNeed() may block on an mmap advise / I/O + // hint and we don't want to serialize concurrent Reads on it. + st = file->WillNeed(ranges); // As this is optimisation only, I/O failures should not be treated as fatal if (st.IsIOError()) { return Status::OK(); @@ -196,22 +212,28 @@ struct ReadRangeCache::Impl { return st; } - // Read the given range from the cache, blocking if needed. Cannot read a range - // that spans cache entries. - virtual Result> Read(ReadRange range) { + // Read the given range from the cache, blocking if needed. Cannot read a + // range that spans cache entries. + Result> Read(ReadRange range) { if (range.length == 0) { static const uint8_t byte = 0; return std::make_shared(&byte, 0); } - const auto it = std::lower_bound( - entries.begin(), entries.end(), range, - [](const RangeCacheEntry& entry, const ReadRange& range) { - return entry.range.offset + entry.range.length < range.offset + range.length; - }); - if (it != entries.end() && it->range.Contains(range)) { - auto fut = MaybeRead(&*it); - ARROW_ASSIGN_OR_RAISE(auto buf, fut.result()); + Future> fut; + int64_t slice_offset = 0; + { + std::unique_lock guard(entry_mutex); + const auto it = std::lower_bound( + entries.begin(), entries.end(), range, + [](const RangeCacheEntry& entry, const ReadRange& range) { + return entry.range.offset + entry.range.length < range.offset + range.length; + }); + if (it == entries.end() || !it->range.Contains(range)) { + return Status::Invalid("ReadRangeCache did not find matching cache entry"); + } + fut = MaybeRead(&*it); + slice_offset = range.offset - it->range.offset; if (options.lazy && options.prefetch_limit > 0) { int64_t num_prefetched = 0; for (auto next_it = it + 1; @@ -224,53 +246,94 @@ struct ReadRangeCache::Impl { ++num_prefetched; } } - return SliceBuffer(std::move(buf), range.offset - it->range.offset, range.length); } - return Status::Invalid("ReadRangeCache did not find matching cache entry"); + // Drop the lock before blocking on the I/O future so other threads can + // still do lookups while a previously queued read is in flight. + ARROW_ASSIGN_OR_RAISE(auto buf, fut.result()); + return SliceBuffer(std::move(buf), slice_offset, range.length); } - virtual Future<> Wait() { + Future<> Wait() { std::vector> futures; - for (auto& entry : entries) { - futures.emplace_back(MaybeRead(&entry)); + { + std::unique_lock guard(entry_mutex); + futures.reserve(entries.size()); + for (auto& entry : entries) { + futures.emplace_back(MaybeRead(&entry)); + } } return AllComplete(futures); } + // Evict cache entries whose byte range is fully contained within + // [start_offset, start_offset + length). + Result EvictEntriesInRange(int64_t start_offset, int64_t length) { + if (length <= 0) { + return 0; + } + const int64_t end_offset = start_offset + length; + int64_t n_evicted = 0; + std::unique_lock guard(entry_mutex); + // entries is sorted by range.offset, so we can fast-forward to the first + // entry whose offset is >= start_offset and stop once we pass end_offset. + auto it = std::lower_bound( + entries.begin(), entries.end(), start_offset, + [](const RangeCacheEntry& entry, int64_t offset) { + return entry.range.offset < offset; + }); + while (it != entries.end() && it->range.offset < end_offset) { + if (it->range.offset + it->range.length <= end_offset) { + it = entries.erase(it); + ++n_evicted; + } else { + // Entry extends beyond the requested window (e.g. coalesced with a + // neighboring range). Leave it alone and keep scanning in case + // smaller siblings follow. + ++it; + } + } + return n_evicted; + } + // Return a Future that completes when the given ranges have been read. - virtual Future<> WaitFor(std::vector ranges) { + Future<> WaitFor(std::vector ranges) { auto end = std::remove_if(ranges.begin(), ranges.end(), [](const ReadRange& range) { return range.length == 0; }); ranges.resize(end - ranges.begin()); std::vector> futures; futures.reserve(ranges.size()); - for (auto& range : ranges) { - const auto it = std::lower_bound( - entries.begin(), entries.end(), range, - [](const RangeCacheEntry& entry, const ReadRange& range) { - return entry.range.offset + entry.range.length < range.offset + range.length; - }); - if (it != entries.end() && it->range.Contains(range)) { - futures.push_back(Future<>(MaybeRead(&*it))); - } else { - return Status::Invalid("Range was not requested for caching: offset=", - range.offset, " length=", range.length); + { + std::unique_lock guard(entry_mutex); + for (auto& range : ranges) { + const auto it = std::lower_bound( + entries.begin(), entries.end(), range, + [](const RangeCacheEntry& entry, const ReadRange& range) { + return entry.range.offset + entry.range.length < + range.offset + range.length; + }); + if (it != entries.end() && it->range.Contains(range)) { + futures.push_back(Future<>(MaybeRead(&*it))); + } else { + return Status::Invalid("Range was not requested for caching: offset=", + range.offset, " length=", range.length); + } } } return AllComplete(futures); } }; -// Don't read ranges when they're first added. Instead, wait until they're requested -// (either through Read or WaitFor). +// Don't read ranges when they're first added. Instead, wait until they're +// requested (either through Read or WaitFor). Thread safety is inherited from +// the base Impl: both MakeCacheEntries and MaybeRead are only ever invoked +// from the public entry points, and those hold the base class's entry_mutex +// before delegating. struct ReadRangeCache::LazyImpl : public ReadRangeCache::Impl { - // Protect against concurrent modification of entries[i]->future - std::mutex entry_mutex; - virtual ~LazyImpl() = default; Future> MaybeRead(RangeCacheEntry* entry) override { - // Called by superclass Read()/WaitFor() so we have the lock + // Called by the superclass under entry_mutex, so it is safe to mutate + // entry->future here. if (!entry->future.is_valid()) { entry->future = file->ReadAsync(ctx, entry->range.offset, entry->range.length); } @@ -288,26 +351,6 @@ struct ReadRangeCache::LazyImpl : public ReadRangeCache::Impl { } return new_entries; } - - Status Cache(std::vector ranges) override { - std::unique_lock guard(entry_mutex); - return ReadRangeCache::Impl::Cache(std::move(ranges)); - } - - Result> Read(ReadRange range) override { - std::unique_lock guard(entry_mutex); - return ReadRangeCache::Impl::Read(range); - } - - Future<> Wait() override { - std::unique_lock guard(entry_mutex); - return ReadRangeCache::Impl::Wait(); - } - - Future<> WaitFor(std::vector ranges) override { - std::unique_lock guard(entry_mutex); - return ReadRangeCache::Impl::WaitFor(std::move(ranges)); - } }; ReadRangeCache::ReadRangeCache(std::shared_ptr owned_file, @@ -336,6 +379,10 @@ Future<> ReadRangeCache::WaitFor(std::vector ranges) { return impl_->WaitFor(std::move(ranges)); } +Result ReadRangeCache::EvictEntriesInRange(int64_t start_offset, int64_t length) { + return impl_->EvictEntriesInRange(start_offset, length); +} + } // namespace internal } // namespace io } // namespace arrow diff --git a/cpp/src/arrow/io/caching.h b/cpp/src/arrow/io/caching.h index e2b911fafdbb..97671d6b2b70 100644 --- a/cpp/src/arrow/io/caching.h +++ b/cpp/src/arrow/io/caching.h @@ -142,6 +142,21 @@ class ARROW_EXPORT ReadRangeCache { /// \brief Wait until all given ranges have been cached. Future<> WaitFor(std::vector ranges); + /// \brief Evict cache entries whose byte range is fully contained within + /// [start_offset, start_offset + length). + /// + /// This releases the memory held by those entries, allowing buffers to be + /// freed as soon as no other owner retains a reference. Entries that are not + /// fully contained in the given window (for example, because I/O range + /// coalescing merged them with an adjacent region the caller still needs) + /// are not evicted, so this method is safe to call even if some cached + /// ranges have been merged across unrelated consumers. + /// + /// \param[in] start_offset Start of the window (inclusive). + /// \param[in] length Length of the window. + /// \return Number of cache entries that were evicted. + Result EvictEntriesInRange(int64_t start_offset, int64_t length); + protected: struct Impl; struct LazyImpl; diff --git a/cpp/src/arrow/io/memory_test.cc b/cpp/src/arrow/io/memory_test.cc index 1b2c7bdbf393..812a4d133307 100644 --- a/cpp/src/arrow/io/memory_test.cc +++ b/cpp/src/arrow/io/memory_test.cc @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +#include #include #include #include @@ -25,6 +26,7 @@ #include #include #include +#include #include #include @@ -917,6 +919,186 @@ TEST(RangeReadCache, LazyWithPrefetching) { ASSERT_RAISES(Invalid, cache.Read({25, 2})); } +TEST(RangeReadCache, EvictEntriesInRange) { + // GH-39808: entries cached by PreBuffer()-style code need to be evictable so + // that memory usage remains bounded while iterating over a large Parquet + // file. + std::string data = "abcdefghijklmnopqrstuvwxyz"; + + for (auto lazy : std::vector{false, true}) { + SCOPED_TRACE(lazy); + CacheOptions options = CacheOptions::Defaults(); + // Disable coalescing so each requested range becomes its own entry. + options.hole_size_limit = 0; + options.range_size_limit = 10; + options.lazy = lazy; + + auto file = std::make_shared(std::make_shared(data)); + internal::ReadRangeCache cache(file, {}, options); + + // Cache three disjoint ranges, each forming its own entry. + ASSERT_OK(cache.Cache({{1, 2}, {10, 4}, {20, 2}})); + + // Sanity: all three ranges are readable. + ASSERT_OK_AND_ASSIGN(auto buf, cache.Read({1, 2})); + AssertBufferEqual(*buf, "bc"); + ASSERT_OK_AND_ASSIGN(buf, cache.Read({10, 4})); + AssertBufferEqual(*buf, "klmn"); + ASSERT_OK_AND_ASSIGN(buf, cache.Read({20, 2})); + AssertBufferEqual(*buf, "uv"); + + // A window that doesn't fully contain any entry evicts nothing. + ASSERT_OK_AND_ASSIGN(int64_t evicted, cache.EvictEntriesInRange(0, 1)); + ASSERT_EQ(0, evicted); + ASSERT_OK_AND_ASSIGN(evicted, cache.EvictEntriesInRange(11, 2)); + ASSERT_EQ(0, evicted); + + // A zero- or negative-length window is a no-op. + ASSERT_OK_AND_ASSIGN(evicted, cache.EvictEntriesInRange(0, 0)); + ASSERT_EQ(0, evicted); + ASSERT_OK_AND_ASSIGN(evicted, cache.EvictEntriesInRange(10, -5)); + ASSERT_EQ(0, evicted); + + // Windows that partially overlap with an entry but don't fully contain + // it must also leave the entry in place. + ASSERT_OK_AND_ASSIGN(evicted, cache.EvictEntriesInRange(0, 2)); + ASSERT_EQ(0, evicted); + ASSERT_OK_AND_ASSIGN(evicted, cache.EvictEntriesInRange(10, 3)); + ASSERT_EQ(0, evicted); + + // Evicting the exact range of the middle entry removes only that entry. + ASSERT_OK_AND_ASSIGN(evicted, cache.EvictEntriesInRange(10, 4)); + ASSERT_EQ(1, evicted); + // Other entries are still readable. + ASSERT_OK_AND_ASSIGN(buf, cache.Read({1, 2})); + AssertBufferEqual(*buf, "bc"); + ASSERT_OK_AND_ASSIGN(buf, cache.Read({20, 2})); + AssertBufferEqual(*buf, "uv"); + // Evicted entry is gone. + ASSERT_RAISES(Invalid, cache.Read({10, 4})); + + // A wider window evicts every remaining fully-contained entry in one call. + ASSERT_OK_AND_ASSIGN(evicted, cache.EvictEntriesInRange(0, 100)); + ASSERT_EQ(2, evicted); + ASSERT_RAISES(Invalid, cache.Read({1, 2})); + ASSERT_RAISES(Invalid, cache.Read({20, 2})); + + // Evicting an already-empty cache is a safe no-op. + ASSERT_OK_AND_ASSIGN(evicted, cache.EvictEntriesInRange(0, 100)); + ASSERT_EQ(0, evicted); + } +} + +TEST(RangeReadCache, ConcurrentReadAndEvict) { + // GH-39808: the Parquet dataset scanner calls EvictEntriesInRange from the + // thread-pool continuation that runs after a row group is decoded, while + // other threads may still be calling Read() for column chunks of other + // in-flight row groups. Exercise that pattern explicitly by slamming the + // cache with parallel Read()s interleaved with Evict()s and make sure we + // don't hit UB (iterator invalidation, torn reads, etc.). + constexpr int kNumRanges = 64; + constexpr int kRangeSize = 64; + constexpr int kIterations = 50; + std::string data(kNumRanges * kRangeSize, 'x'); + + for (auto lazy : std::vector{false, true}) { + SCOPED_TRACE(lazy); + CacheOptions options = CacheOptions::Defaults(); + // No coalescing: each range is its own entry so we can evict them + // individually without fighting the coalescing heuristic. + options.hole_size_limit = 0; + options.range_size_limit = kRangeSize; + options.lazy = lazy; + + auto file = std::make_shared(std::make_shared(data)); + internal::ReadRangeCache cache(file, {}, options); + + std::vector ranges; + ranges.reserve(kNumRanges); + for (int i = 0; i < kNumRanges; ++i) { + ranges.push_back({static_cast(i * kRangeSize), kRangeSize}); + } + ASSERT_OK(cache.Cache(ranges)); + + // Half of the threads repeatedly read the upper half of the ranges. + // The other half repeatedly evict and re-cache the lower half. Under + // the old code this would race on the shared `entries` vector. + std::atomic stop{false}; + std::atomic failures{0}; + + auto reader_fn = [&]() { + while (!stop.load(std::memory_order_relaxed)) { + for (int i = kNumRanges / 2; i < kNumRanges; ++i) { + auto result = cache.Read(ranges[i]); + if (!result.ok() || (*result)->size() != kRangeSize) { + failures.fetch_add(1, std::memory_order_relaxed); + } + } + } + }; + auto evictor_fn = [&]() { + for (int iter = 0; iter < kIterations; ++iter) { + // Evict every lower-half entry. + for (int i = 0; i < kNumRanges / 2; ++i) { + auto result = cache.EvictEntriesInRange(ranges[i].offset, ranges[i].length); + if (!result.ok()) { + failures.fetch_add(1, std::memory_order_relaxed); + } + } + // Re-cache them so the next iteration has something to evict. + std::vector lower(ranges.begin(), ranges.begin() + kNumRanges / 2); + if (!cache.Cache(lower).ok()) { + failures.fetch_add(1, std::memory_order_relaxed); + } + } + stop.store(true, std::memory_order_relaxed); + }; + + std::vector threads; + for (int i = 0; i < 4; ++i) threads.emplace_back(reader_fn); + threads.emplace_back(evictor_fn); + for (auto& t : threads) t.join(); + ASSERT_EQ(0, failures.load()); + + // Every upper-half range is still readable after the torture loop. + for (int i = kNumRanges / 2; i < kNumRanges; ++i) { + ASSERT_OK_AND_ASSIGN(auto buf, cache.Read(ranges[i])); + ASSERT_EQ(kRangeSize, buf->size()); + } + } +} + +TEST(RangeReadCache, EvictEntriesInRangeSpanningEntry) { + // Coalesced entries (entries that were merged with a neighboring range at + // Cache() time) must not be evicted unless the eviction window fully + // contains them, otherwise we would drop bytes the caller still needs. + std::string data(40, 'x'); + + CacheOptions options = CacheOptions::Defaults(); + // Large hole_size_limit forces coalescing of adjacent ranges into one entry. + options.hole_size_limit = 100; + options.range_size_limit = 200; + + auto file = std::make_shared(std::make_shared(data)); + internal::ReadRangeCache cache(file, {}, options); + + // These two ranges get coalesced into a single entry [1, 14). + ASSERT_OK(cache.Cache({{1, 3}, {10, 4}})); + + // Trying to evict only one of the two "logical" ranges must not drop the + // coalesced entry - it still holds data the other logical range needs. + ASSERT_OK_AND_ASSIGN(int64_t evicted, cache.EvictEntriesInRange(1, 3)); + ASSERT_EQ(0, evicted); + ASSERT_OK_AND_ASSIGN(auto buf, cache.Read({10, 4})); + ASSERT_EQ(4, buf->size()); + + // A wide window that contains the whole coalesced entry evicts it. + ASSERT_OK_AND_ASSIGN(evicted, cache.EvictEntriesInRange(0, 20)); + ASSERT_EQ(1, evicted); + ASSERT_RAISES(Invalid, cache.Read({1, 3})); + ASSERT_RAISES(Invalid, cache.Read({10, 4})); +} + TEST(CacheOptions, Basics) { auto check = [](const CacheOptions actual, const double expected_hole_size_limit_MiB, const double expected_range_size_limit_MiB) -> void { diff --git a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc index e2384972cf55..b19a7d1713e7 100644 --- a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc +++ b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc @@ -2670,6 +2670,116 @@ TEST(TestArrowReadWrite, GetRecordBatchReaderNoColumns) { ASSERT_EQ(actual_batch->num_rows(), num_rows); } +// GH-39808: bytes cached by PreBuffer() for an already-decoded row group +// should be released when the caller explicitly evicts them, otherwise +// Dataset.to_batches accumulates memory across the lifetime of an open +// FileReader. +TEST(TestArrowReadWrite, EvictPreBufferedData) { + ArrowReaderProperties properties = default_arrow_reader_properties(); + properties.set_pre_buffer(true); + const int num_rows = 1024; + const int row_group_size = 256; + const int num_columns = 3; + const std::vector row_groups = {0, 1, 2, 3}; + const std::vector column_indices = {0, 1, 2}; + + std::shared_ptr table; + ASSERT_NO_FATAL_FAILURE(MakeDoubleTable(num_columns, num_rows, 1, &table)); + + std::shared_ptr buffer; + ASSERT_NO_FATAL_FAILURE(WriteTableToBuffer(table, row_group_size, + default_arrow_writer_properties(), &buffer)); + + std::unique_ptr reader; + FileReaderBuilder builder; + ASSERT_OK(builder.Open(std::make_shared(buffer))); + ASSERT_OK(builder.properties(properties)->Build(&reader)); + ASSERT_EQ(reader->num_row_groups(), static_cast(row_groups.size())); + + // Pre-buffer every row group/column so the cache is fully populated. + reader->parquet_reader()->PreBuffer(row_groups, column_indices, + ::arrow::io::IOContext(), + ::arrow::io::CacheOptions::LazyDefaults()); + ASSERT_OK( + reader->parquet_reader()->WhenBuffered(row_groups, column_indices).status()); + + // Decode the first row group; reads go through the cache. + std::shared_ptr
rg_table; + ASSERT_OK(reader->ReadRowGroup(/*i=*/0, column_indices, &rg_table)); + ASSERT_EQ(rg_table->num_rows(), row_group_size); + + // Evict only row group 0. The ranges for the other row groups are untouched, + // so they must still be readable. + reader->parquet_reader()->EvictPreBufferedData({0}, column_indices); + ASSERT_OK(reader->ReadRowGroup(/*i=*/1, column_indices, &rg_table)); + ASSERT_EQ(rg_table->num_rows(), row_group_size); + + // Evicting twice is a no-op (and must not crash). + reader->parquet_reader()->EvictPreBufferedData({0}, column_indices); + + // Calling EvictPreBufferedData on a reader that never called PreBuffer is + // also a no-op. + std::unique_ptr no_prebuffer_reader; + FileReaderBuilder no_prebuffer_builder; + ASSERT_OK(no_prebuffer_builder.Open(std::make_shared(buffer))); + ASSERT_OK(no_prebuffer_builder.Build(&no_prebuffer_reader)); + no_prebuffer_reader->parquet_reader()->EvictPreBufferedData(row_groups, + column_indices); +} + +// GH-39808: when Dataset.to_batches-style iteration drives the async +// RecordBatchGenerator, each row group's pre-buffered bytes should be +// released as soon as the row group has been converted into record batches, +// so the overall memory footprint is independent of how many row groups the +// file contains. +TEST(TestArrowReadWrite, GetRecordBatchGeneratorReleasesPreBufferedRowGroups) { + ArrowReaderProperties properties = default_arrow_reader_properties(); + properties.set_pre_buffer(true); + // Read one row group at a time so the test deterministically exercises the + // per-row-group eviction path. + properties.set_batch_size(256); + + const int num_rows = 1024; + const int row_group_size = 256; + const int num_columns = 2; + + std::shared_ptr
table; + ASSERT_NO_FATAL_FAILURE(MakeDoubleTable(num_columns, num_rows, 1, &table)); + + std::shared_ptr buffer; + ASSERT_NO_FATAL_FAILURE(WriteTableToBuffer(table, row_group_size, + default_arrow_writer_properties(), &buffer)); + + std::shared_ptr reader; + { + std::unique_ptr unique_reader; + FileReaderBuilder builder; + ASSERT_OK(builder.Open(std::make_shared(buffer))); + ASSERT_OK(builder.properties(properties)->Build(&unique_reader)); + reader = std::move(unique_reader); + } + ASSERT_EQ(reader->num_row_groups(), num_rows / row_group_size); + + // Drive the generator exactly as ScanBatchesAsync does. + ASSERT_OK_AND_ASSIGN(auto batch_generator, + reader->GetRecordBatchGenerator(reader, {0, 1, 2, 3}, {0, 1})); + std::vector> batches; + for (int i = 0; i < reader->num_row_groups(); ++i) { + auto fut = batch_generator(); + ASSERT_OK_AND_ASSIGN(auto batch, fut.result()); + ASSERT_NE(batch, nullptr); + batches.push_back(std::move(batch)); + } + // Generator is drained. + auto fut_end = batch_generator(); + ASSERT_OK_AND_ASSIGN(auto end_batch, fut_end.result()); + ASSERT_EQ(end_batch, nullptr); + + ASSERT_OK_AND_ASSIGN( + auto actual, ::arrow::Table::FromRecordBatches(batches[0]->schema(), batches)); + AssertTablesEqual(*table, *actual, /*same_chunk_layout=*/false); +} + TEST(TestArrowReadWrite, GetRecordBatchGenerator) { ArrowReaderProperties properties = default_arrow_reader_properties(); const int num_rows = 1024; diff --git a/cpp/src/parquet/arrow/reader.cc b/cpp/src/parquet/arrow/reader.cc index a60af69aec9f..5c436b665ff2 100644 --- a/cpp/src/parquet/arrow/reader.cc +++ b/cpp/src/parquet/arrow/reader.cc @@ -1189,12 +1189,22 @@ class RowGroupGenerator { const int row_group, const std::vector& column_indices) { // Skips bound checks/pre-buffering, since we've done that already const int64_t batch_size = self->properties().batch_size(); + const bool pre_buffered = self->properties().pre_buffer(); return self->DecodeRowGroups(self, {row_group}, column_indices, cpu_executor) - .Then([batch_size](const std::shared_ptr
& table) + .Then([batch_size, self, row_group, column_indices = column_indices, + pre_buffered](const std::shared_ptr
& table) -> ::arrow::Result { ::arrow::TableBatchReader table_reader(*table); table_reader.set_chunksize(batch_size); ARROW_ASSIGN_OR_RAISE(auto batches, table_reader.ToRecordBatches()); + // GH-39808: once a row group has been fully decoded into Arrow + // arrays, the bytes cached by PreBuffer() for that row group are no + // longer needed. Release them so memory usage stays bounded while + // iterating a large dataset instead of growing until the whole file + // has been read. + if (pre_buffered) { + self->parquet_reader()->EvictPreBufferedData({row_group}, column_indices); + } return ::arrow::MakeVectorGenerator(std::move(batches)); }); } diff --git a/cpp/src/parquet/file_reader.cc b/cpp/src/parquet/file_reader.cc index af7ccfd7ad7d..ad9652ee15c9 100644 --- a/cpp/src/parquet/file_reader.cc +++ b/cpp/src/parquet/file_reader.cc @@ -20,6 +20,7 @@ #include #include #include +#include #include #include #include @@ -432,6 +433,41 @@ class SerializedFile : public ParquetFileReader::Contents { return cached_source_->WaitFor(ranges); } + // Evict cached bytes that were populated by PreBuffer() for the given row + // groups and column indices. Callers should only invoke this once the + // corresponding row group data has been fully decoded and no readers are + // holding a reference to the cached buffers. + void EvictPreBufferedData(const std::vector& row_groups, + const std::vector& column_indices) { + if (!cached_source_) { + return; + } + for (int row : row_groups) { + if (column_indices.empty()) { + continue; + } + // Bounding box of the row group's column chunk ranges. Using the bounding + // box (instead of per-column ranges) allows the cache to evict coalesced + // entries that cover multiple columns of the same row group, while + // leaving alone any entry that may have merged across row groups (in + // which case the merged entry extends beyond this bounding box and is + // not fully contained). + int64_t min_start = std::numeric_limits::max(); + int64_t max_end = std::numeric_limits::min(); + for (int col : column_indices) { + auto range = + ComputeColumnChunkRange(file_metadata_.get(), source_size_, row, col); + min_start = std::min(min_start, range.offset); + max_end = std::max(max_end, range.offset + range.length); + } + if (max_end > min_start) { + PARQUET_THROW_NOT_OK( + cached_source_->EvictEntriesInRange(min_start, max_end - min_start) + .status()); + } + } + } + // Metadata/footer parsing. Divided up to separate sync/async paths, and to use // exceptions for error handling (with the async path converting to Future/Status). @@ -905,6 +941,14 @@ void ParquetFileReader::PreBuffer(const std::vector& row_groups, file->PreBuffer(row_groups, column_indices, ctx, options); } +void ParquetFileReader::EvictPreBufferedData( + const std::vector& row_groups, const std::vector& column_indices) { + // Access private methods here + SerializedFile* file = + ::arrow::internal::checked_cast(contents_.get()); + file->EvictPreBufferedData(row_groups, column_indices); +} + Result> ParquetFileReader::GetReadRanges( const std::vector& row_groups, const std::vector& column_indices, int64_t hole_size_limit, int64_t range_size_limit) { diff --git a/cpp/src/parquet/file_reader.h b/cpp/src/parquet/file_reader.h index c42163276cda..db52df9c7820 100644 --- a/cpp/src/parquet/file_reader.h +++ b/cpp/src/parquet/file_reader.h @@ -201,6 +201,23 @@ class PARQUET_EXPORT ParquetFileReader { const ::arrow::io::IOContext& ctx, const ::arrow::io::CacheOptions& options); + /// Release bytes buffered by a previous call to PreBuffer() for the + /// specified row groups and column indices. + /// + /// This releases the memory held by those cached ranges, allowing the + /// underlying buffers to be freed as soon as no other owner retains a + /// reference. Callers must ensure that the targeted row groups have been + /// fully decoded before invoking this method; otherwise subsequent reads + /// of the same row groups may fail because the cached bytes are gone. + /// + /// It is safe to call this concurrently with reads of other row groups: + /// cache entries that were coalesced across the targeted row groups and + /// unrelated data are left intact. + /// + /// Calling this method when PreBuffer() has not been called is a no-op. + void EvictPreBufferedData(const std::vector& row_groups, + const std::vector& column_indices); + /// Retrieve the list of byte ranges that would need to be read to retrieve /// the data for the specified row groups and column indices. /// From 43dd3f13eb72eacaec6502baff6f21d9b5c62506 Mon Sep 17 00:00:00 2001 From: Justin Li Date: Sun, 14 Jun 2026 22:30:51 -0400 Subject: [PATCH 02/12] GH-39808: [C++][Parquet] Apply clang-format to fix lint Reformat the pre-buffer eviction changes with clang-format 18.1.8 to satisfy the CI lint job. Whitespace and line-wrapping only; no behavior change. --- cpp/src/arrow/io/caching.cc | 28 +++++++++---------- .../parquet/arrow/arrow_reader_writer_test.cc | 10 +++---- cpp/src/parquet/file_reader.cc | 7 ++--- 3 files changed, 21 insertions(+), 24 deletions(-) diff --git a/cpp/src/arrow/io/caching.cc b/cpp/src/arrow/io/caching.cc index abb1cafcb059..ad13f192498d 100644 --- a/cpp/src/arrow/io/caching.cc +++ b/cpp/src/arrow/io/caching.cc @@ -194,8 +194,8 @@ struct ReadRangeCache::Impl { // Add new entries, themselves ordered by offset if (entries.size() > 0) { std::vector merged(entries.size() + new_entries.size()); - std::merge(entries.begin(), entries.end(), new_entries.begin(), - new_entries.end(), merged.begin()); + std::merge(entries.begin(), entries.end(), new_entries.begin(), new_entries.end(), + merged.begin()); entries = std::move(merged); } else { entries = std::move(new_entries); @@ -276,11 +276,10 @@ struct ReadRangeCache::Impl { std::unique_lock guard(entry_mutex); // entries is sorted by range.offset, so we can fast-forward to the first // entry whose offset is >= start_offset and stop once we pass end_offset. - auto it = std::lower_bound( - entries.begin(), entries.end(), start_offset, - [](const RangeCacheEntry& entry, int64_t offset) { - return entry.range.offset < offset; - }); + auto it = std::lower_bound(entries.begin(), entries.end(), start_offset, + [](const RangeCacheEntry& entry, int64_t offset) { + return entry.range.offset < offset; + }); while (it != entries.end() && it->range.offset < end_offset) { if (it->range.offset + it->range.length <= end_offset) { it = entries.erase(it); @@ -305,12 +304,12 @@ struct ReadRangeCache::Impl { { std::unique_lock guard(entry_mutex); for (auto& range : ranges) { - const auto it = std::lower_bound( - entries.begin(), entries.end(), range, - [](const RangeCacheEntry& entry, const ReadRange& range) { - return entry.range.offset + entry.range.length < - range.offset + range.length; - }); + const auto it = + std::lower_bound(entries.begin(), entries.end(), range, + [](const RangeCacheEntry& entry, const ReadRange& range) { + return entry.range.offset + entry.range.length < + range.offset + range.length; + }); if (it != entries.end() && it->range.Contains(range)) { futures.push_back(Future<>(MaybeRead(&*it))); } else { @@ -379,7 +378,8 @@ Future<> ReadRangeCache::WaitFor(std::vector ranges) { return impl_->WaitFor(std::move(ranges)); } -Result ReadRangeCache::EvictEntriesInRange(int64_t start_offset, int64_t length) { +Result ReadRangeCache::EvictEntriesInRange(int64_t start_offset, + int64_t length) { return impl_->EvictEntriesInRange(start_offset, length); } diff --git a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc index b19a7d1713e7..aed490041b34 100644 --- a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc +++ b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc @@ -2700,8 +2700,7 @@ TEST(TestArrowReadWrite, EvictPreBufferedData) { reader->parquet_reader()->PreBuffer(row_groups, column_indices, ::arrow::io::IOContext(), ::arrow::io::CacheOptions::LazyDefaults()); - ASSERT_OK( - reader->parquet_reader()->WhenBuffered(row_groups, column_indices).status()); + ASSERT_OK(reader->parquet_reader()->WhenBuffered(row_groups, column_indices).status()); // Decode the first row group; reads go through the cache. std::shared_ptr
rg_table; @@ -2723,8 +2722,7 @@ TEST(TestArrowReadWrite, EvictPreBufferedData) { FileReaderBuilder no_prebuffer_builder; ASSERT_OK(no_prebuffer_builder.Open(std::make_shared(buffer))); ASSERT_OK(no_prebuffer_builder.Build(&no_prebuffer_reader)); - no_prebuffer_reader->parquet_reader()->EvictPreBufferedData(row_groups, - column_indices); + no_prebuffer_reader->parquet_reader()->EvictPreBufferedData(row_groups, column_indices); } // GH-39808: when Dataset.to_batches-style iteration drives the async @@ -2775,8 +2773,8 @@ TEST(TestArrowReadWrite, GetRecordBatchGeneratorReleasesPreBufferedRowGroups) { ASSERT_OK_AND_ASSIGN(auto end_batch, fut_end.result()); ASSERT_EQ(end_batch, nullptr); - ASSERT_OK_AND_ASSIGN( - auto actual, ::arrow::Table::FromRecordBatches(batches[0]->schema(), batches)); + ASSERT_OK_AND_ASSIGN(auto actual, + ::arrow::Table::FromRecordBatches(batches[0]->schema(), batches)); AssertTablesEqual(*table, *actual, /*same_chunk_layout=*/false); } diff --git a/cpp/src/parquet/file_reader.cc b/cpp/src/parquet/file_reader.cc index ad9652ee15c9..1381c578da24 100644 --- a/cpp/src/parquet/file_reader.cc +++ b/cpp/src/parquet/file_reader.cc @@ -462,8 +462,7 @@ class SerializedFile : public ParquetFileReader::Contents { } if (max_end > min_start) { PARQUET_THROW_NOT_OK( - cached_source_->EvictEntriesInRange(min_start, max_end - min_start) - .status()); + cached_source_->EvictEntriesInRange(min_start, max_end - min_start).status()); } } } @@ -941,8 +940,8 @@ void ParquetFileReader::PreBuffer(const std::vector& row_groups, file->PreBuffer(row_groups, column_indices, ctx, options); } -void ParquetFileReader::EvictPreBufferedData( - const std::vector& row_groups, const std::vector& column_indices) { +void ParquetFileReader::EvictPreBufferedData(const std::vector& row_groups, + const std::vector& column_indices) { // Access private methods here SerializedFile* file = ::arrow::internal::checked_cast(contents_.get()); From 008e5b771cef900f7961054d5825cac4dc92b018 Mon Sep 17 00:00:00 2001 From: Justin Li Date: Tue, 16 Jun 2026 03:18:33 -0400 Subject: [PATCH 03/12] GH-39808: [C++][Parquet] Return evicted-entry count from EvictPreBufferedData --- cpp/src/parquet/file_reader.cc | 23 ++++++++++++++--------- cpp/src/parquet/file_reader.h | 4 ++-- 2 files changed, 16 insertions(+), 11 deletions(-) diff --git a/cpp/src/parquet/file_reader.cc b/cpp/src/parquet/file_reader.cc index 1381c578da24..dd816f26c5ce 100644 --- a/cpp/src/parquet/file_reader.cc +++ b/cpp/src/parquet/file_reader.cc @@ -436,12 +436,14 @@ class SerializedFile : public ParquetFileReader::Contents { // Evict cached bytes that were populated by PreBuffer() for the given row // groups and column indices. Callers should only invoke this once the // corresponding row group data has been fully decoded and no readers are - // holding a reference to the cached buffers. - void EvictPreBufferedData(const std::vector& row_groups, - const std::vector& column_indices) { + // holding a reference to the cached buffers. Returns the number of cache + // entries that were evicted. + int64_t EvictPreBufferedData(const std::vector& row_groups, + const std::vector& column_indices) { if (!cached_source_) { - return; + return 0; } + int64_t total_evicted = 0; for (int row : row_groups) { if (column_indices.empty()) { continue; @@ -461,10 +463,13 @@ class SerializedFile : public ParquetFileReader::Contents { max_end = std::max(max_end, range.offset + range.length); } if (max_end > min_start) { - PARQUET_THROW_NOT_OK( - cached_source_->EvictEntriesInRange(min_start, max_end - min_start).status()); + PARQUET_ASSIGN_OR_THROW( + int64_t evicted, + cached_source_->EvictEntriesInRange(min_start, max_end - min_start)); + total_evicted += evicted; } } + return total_evicted; } // Metadata/footer parsing. Divided up to separate sync/async paths, and to use @@ -940,12 +945,12 @@ void ParquetFileReader::PreBuffer(const std::vector& row_groups, file->PreBuffer(row_groups, column_indices, ctx, options); } -void ParquetFileReader::EvictPreBufferedData(const std::vector& row_groups, - const std::vector& column_indices) { +int64_t ParquetFileReader::EvictPreBufferedData(const std::vector& row_groups, + const std::vector& column_indices) { // Access private methods here SerializedFile* file = ::arrow::internal::checked_cast(contents_.get()); - file->EvictPreBufferedData(row_groups, column_indices); + return file->EvictPreBufferedData(row_groups, column_indices); } Result> ParquetFileReader::GetReadRanges( diff --git a/cpp/src/parquet/file_reader.h b/cpp/src/parquet/file_reader.h index db52df9c7820..fad5ee9eae38 100644 --- a/cpp/src/parquet/file_reader.h +++ b/cpp/src/parquet/file_reader.h @@ -215,8 +215,8 @@ class PARQUET_EXPORT ParquetFileReader { /// unrelated data are left intact. /// /// Calling this method when PreBuffer() has not been called is a no-op. - void EvictPreBufferedData(const std::vector& row_groups, - const std::vector& column_indices); + int64_t EvictPreBufferedData(const std::vector& row_groups, + const std::vector& column_indices); /// Retrieve the list of byte ranges that would need to be read to retrieve /// the data for the specified row groups and column indices. From ac1d90f7051095a34b6197d7a5d49e1897037433 Mon Sep 17 00:00:00 2001 From: Justin Li Date: Tue, 16 Jun 2026 03:47:46 -0400 Subject: [PATCH 04/12] GH-39808: [C++][Parquet] Evict coalesced cache entries spanning row groups --- .../parquet/arrow/arrow_reader_writer_test.cc | 91 +++++++++++++++++ cpp/src/parquet/file_reader.cc | 97 ++++++++++++++----- 2 files changed, 164 insertions(+), 24 deletions(-) diff --git a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc index aed490041b34..5071325ba3ab 100644 --- a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc +++ b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc @@ -2778,6 +2778,97 @@ TEST(TestArrowReadWrite, GetRecordBatchGeneratorReleasesPreBufferedRowGroups) { AssertTablesEqual(*table, *actual, /*same_chunk_layout=*/false); } +// GH-39808: with default coalescing a single cache entry can span adjacent row +// groups. Per-row-group eviction must still release such an entry, but only +// once every row group it covers has been evicted. +TEST(TestArrowReadWrite, EvictPreBufferedDataReleasesCrossRowGroupEntry) { + ArrowReaderProperties properties = default_arrow_reader_properties(); + properties.set_pre_buffer(true); + const int num_rows = 1024; + const int row_group_size = 256; // 4 row groups + const int num_columns = 2; + const std::vector row_groups = {0, 1, 2, 3}; + const std::vector column_indices = {0, 1}; + + std::shared_ptr
table; + ASSERT_NO_FATAL_FAILURE(MakeDoubleTable(num_columns, num_rows, 1, &table)); + + std::shared_ptr buffer; + ASSERT_NO_FATAL_FAILURE(WriteTableToBuffer(table, row_group_size, + default_arrow_writer_properties(), &buffer)); + + std::unique_ptr reader; + FileReaderBuilder builder; + ASSERT_OK(builder.Open(std::make_shared(buffer))); + ASSERT_OK(builder.properties(properties)->Build(&reader)); + ASSERT_EQ(reader->num_row_groups(), static_cast(row_groups.size())); + + // Force every column chunk of every row group to coalesce into a single + // cache entry that spans all row-group boundaries. + ::arrow::io::CacheOptions options = ::arrow::io::CacheOptions::LazyDefaults(); + options.hole_size_limit = static_cast(buffer->size()); + options.range_size_limit = static_cast(buffer->size()); + reader->parquet_reader()->PreBuffer(row_groups, column_indices, + ::arrow::io::IOContext(), options); + ASSERT_OK(reader->parquet_reader()->WhenBuffered(row_groups, column_indices).status()); + + // Evicting any strict subset of the row groups leaves the spanning entry in + // place: it is not fully contained in any run that omits a row group it + // still covers, so each of these calls evicts nothing. + ASSERT_EQ(0, reader->parquet_reader()->EvictPreBufferedData({0}, column_indices)); + ASSERT_EQ(0, reader->parquet_reader()->EvictPreBufferedData({1}, column_indices)); + ASSERT_EQ(0, reader->parquet_reader()->EvictPreBufferedData({2}, column_indices)); + + // Once the final row group completes the contiguous run [0, 3], the spanning + // entry is fully contained and is freed. + ASSERT_EQ(1, reader->parquet_reader()->EvictPreBufferedData({3}, column_indices)); + + // Idempotent: re-evicting frees nothing more. + ASSERT_EQ(0, reader->parquet_reader()->EvictPreBufferedData({3}, column_indices)); +} + +// GH-39808: eviction order is non-deterministic under readahead (row groups +// decode concurrently). The spanning entry must be freed exactly once the gap +// between evicted runs is filled, regardless of order. +TEST(TestArrowReadWrite, EvictPreBufferedDataReleasesCrossRowGroupEntryOutOfOrder) { + ArrowReaderProperties properties = default_arrow_reader_properties(); + properties.set_pre_buffer(true); + const int num_rows = 1024; + const int row_group_size = 256; // 4 row groups + const int num_columns = 2; + const std::vector row_groups = {0, 1, 2, 3}; + const std::vector column_indices = {0, 1}; + + std::shared_ptr
table; + ASSERT_NO_FATAL_FAILURE(MakeDoubleTable(num_columns, num_rows, 1, &table)); + + std::shared_ptr buffer; + ASSERT_NO_FATAL_FAILURE(WriteTableToBuffer(table, row_group_size, + default_arrow_writer_properties(), &buffer)); + + std::unique_ptr reader; + FileReaderBuilder builder; + ASSERT_OK(builder.Open(std::make_shared(buffer))); + ASSERT_OK(builder.properties(properties)->Build(&reader)); + + ::arrow::io::CacheOptions options = ::arrow::io::CacheOptions::LazyDefaults(); + options.hole_size_limit = static_cast(buffer->size()); + options.range_size_limit = static_cast(buffer->size()); + reader->parquet_reader()->PreBuffer(row_groups, column_indices, + ::arrow::io::IOContext(), options); + ASSERT_OK(reader->parquet_reader()->WhenBuffered(row_groups, column_indices).status()); + + // Evict 2, then 0, then 3: runs {2,3} and {0} exist but neither covers the + // whole entry, so nothing is freed. + ASSERT_EQ(0, reader->parquet_reader()->EvictPreBufferedData({2}, column_indices)); + ASSERT_EQ(0, reader->parquet_reader()->EvictPreBufferedData({0}, column_indices)); + ASSERT_EQ(0, reader->parquet_reader()->EvictPreBufferedData({3}, column_indices)); + + // Evicting 1 fills the gap, merging {0} and {2,3} into {0,1,2,3}; the entry + // is now fully contained and freed. + ASSERT_EQ(1, reader->parquet_reader()->EvictPreBufferedData({1}, column_indices)); +} + TEST(TestArrowReadWrite, GetRecordBatchGenerator) { ArrowReaderProperties properties = default_arrow_reader_properties(); const int num_rows = 1024; diff --git a/cpp/src/parquet/file_reader.cc b/cpp/src/parquet/file_reader.cc index dd816f26c5ce..4a67397ee1b5 100644 --- a/cpp/src/parquet/file_reader.cc +++ b/cpp/src/parquet/file_reader.cc @@ -21,7 +21,9 @@ #include #include #include +#include #include +#include #include #include #include @@ -381,6 +383,12 @@ class SerializedFile : public ParquetFileReader::Contents { const ::arrow::io::CacheOptions& options) { cached_source_ = std::make_shared<::arrow::io::internal::ReadRangeCache>(source_, ctx, options); + { + // GH-39808: a fresh cache invalidates any eviction runs we were + // tracking; they referenced the previous cache's byte ranges. + std::lock_guard lock(eviction_mutex_); + evicted_runs_.clear(); + } std::vector<::arrow::io::ReadRange> ranges; prebuffered_column_chunks_.clear(); int num_cols = file_metadata_->num_columns(); @@ -433,41 +441,69 @@ class SerializedFile : public ParquetFileReader::Contents { return cached_source_->WaitFor(ranges); } - // Evict cached bytes that were populated by PreBuffer() for the given row - // groups and column indices. Callers should only invoke this once the - // corresponding row group data has been fully decoded and no readers are - // holding a reference to the cached buffers. Returns the number of cache - // entries that were evicted. + // Evict cached bytes populated by PreBuffer() for the given row groups. I/O + // coalescing can merge column chunks of adjacent row groups into one cache + // entry, so an entry is freed only once every row group it covers has been + // evicted: we track contiguous runs of evicted row groups (row-group index + // order matches byte order) and evict each run's combined byte window. + // Returns the number of cache entries evicted. + // + // Callers must have fully decoded the row groups (no reader still holds the + // cached buffers). Safe to call concurrently for different row groups. int64_t EvictPreBufferedData(const std::vector& row_groups, const std::vector& column_indices) { - if (!cached_source_) { + if (!cached_source_ || column_indices.empty()) { return 0; } int64_t total_evicted = 0; + std::lock_guard lock(eviction_mutex_); for (int row : row_groups) { - if (column_indices.empty()) { - continue; - } - // Bounding box of the row group's column chunk ranges. Using the bounding - // box (instead of per-column ranges) allows the cache to evict coalesced - // entries that cover multiple columns of the same row group, while - // leaving alone any entry that may have merged across row groups (in - // which case the merged entry extends beyond this bounding box and is - // not fully contained). - int64_t min_start = std::numeric_limits::max(); - int64_t max_end = std::numeric_limits::min(); + // Bounding box of this row group's selected column chunks. + int64_t start = std::numeric_limits::max(); + int64_t end = std::numeric_limits::min(); for (int col : column_indices) { auto range = ComputeColumnChunkRange(file_metadata_.get(), source_size_, row, col); - min_start = std::min(min_start, range.offset); - max_end = std::max(max_end, range.offset + range.length); + start = std::min(start, range.offset); + end = std::max(end, range.offset + range.length); + } + if (end <= start) { + continue; + } + + // Skip row groups already covered by a run (idempotent). + auto after = evicted_runs_.upper_bound(row); + if (after != evicted_runs_.begin() && + row <= std::prev(after)->second.last_row_group) { + continue; } - if (max_end > min_start) { - PARQUET_ASSIGN_OR_THROW( - int64_t evicted, - cached_source_->EvictEntriesInRange(min_start, max_end - min_start)); - total_evicted += evicted; + + // Merge with the adjacent runs (one ending at row - 1, one starting at + // row + 1) into a single contiguous run, unioning the byte windows. + int first = row; + int last = row; + auto right = evicted_runs_.find(row + 1); + if (right != evicted_runs_.end()) { + last = right->second.last_row_group; + end = std::max(end, right->second.end_offset); + evicted_runs_.erase(right); } + auto upper = evicted_runs_.lower_bound(row); + if (upper != evicted_runs_.begin()) { + auto left = std::prev(upper); + if (left->second.last_row_group == row - 1) { + first = left->first; + start = std::min(start, left->second.start_offset); + evicted_runs_.erase(left); + } + } + evicted_runs_[first] = EvictedRowGroupRun{last, start, end}; + + // A cross-row-group coalesced entry is fully contained in the run's + // window only once every row group it spans has been evicted. + PARQUET_ASSIGN_OR_THROW(int64_t evicted, + cached_source_->EvictEntriesInRange(start, end - start)); + total_evicted += evicted; } return total_evicted; } @@ -653,6 +689,19 @@ class SerializedFile : public ParquetFileReader::Contents { ReaderProperties properties_; std::shared_ptr page_index_reader_; std::unique_ptr bloom_filter_reader_; + + // GH-39808: contiguous runs of row groups whose pre-buffered bytes have been + // evicted, keyed by each run's first row-group index. A coalesced cache entry + // can span adjacent row groups, so it is freed only once every row group in + // its run has been evicted. Guarded by eviction_mutex_ because row groups are + // decoded (and evicted) concurrently under readahead. + struct EvictedRowGroupRun { + int last_row_group; + int64_t start_offset; + int64_t end_offset; + }; + std::mutex eviction_mutex_; + std::map evicted_runs_; // Maps row group ordinal and prebuffer status of its column chunks in the form of a // bitmap buffer. std::unordered_map> prebuffered_column_chunks_; From 9ac85eb073ad01b6339f711b8c8d0dbc9fa1181a Mon Sep 17 00:00:00 2001 From: Justin Li Date: Tue, 16 Jun 2026 04:00:28 -0400 Subject: [PATCH 05/12] GH-39808: [C++][Parquet] Trim verbose pre-buffer eviction comments --- cpp/src/parquet/arrow/reader.cc | 7 ++----- cpp/src/parquet/file_reader.cc | 2 -- cpp/src/parquet/file_reader.h | 19 ++++++------------- 3 files changed, 8 insertions(+), 20 deletions(-) diff --git a/cpp/src/parquet/arrow/reader.cc b/cpp/src/parquet/arrow/reader.cc index 5c436b665ff2..29a5a2065779 100644 --- a/cpp/src/parquet/arrow/reader.cc +++ b/cpp/src/parquet/arrow/reader.cc @@ -1197,11 +1197,8 @@ class RowGroupGenerator { ::arrow::TableBatchReader table_reader(*table); table_reader.set_chunksize(batch_size); ARROW_ASSIGN_OR_RAISE(auto batches, table_reader.ToRecordBatches()); - // GH-39808: once a row group has been fully decoded into Arrow - // arrays, the bytes cached by PreBuffer() for that row group are no - // longer needed. Release them so memory usage stays bounded while - // iterating a large dataset instead of growing until the whole file - // has been read. + // GH-39808: release this row group's pre-buffered bytes now that it + // is decoded, keeping memory bounded while iterating. if (pre_buffered) { self->parquet_reader()->EvictPreBufferedData({row_group}, column_indices); } diff --git a/cpp/src/parquet/file_reader.cc b/cpp/src/parquet/file_reader.cc index 4a67397ee1b5..aa9bb6501049 100644 --- a/cpp/src/parquet/file_reader.cc +++ b/cpp/src/parquet/file_reader.cc @@ -499,8 +499,6 @@ class SerializedFile : public ParquetFileReader::Contents { } evicted_runs_[first] = EvictedRowGroupRun{last, start, end}; - // A cross-row-group coalesced entry is fully contained in the run's - // window only once every row group it spans has been evicted. PARQUET_ASSIGN_OR_THROW(int64_t evicted, cached_source_->EvictEntriesInRange(start, end - start)); total_evicted += evicted; diff --git a/cpp/src/parquet/file_reader.h b/cpp/src/parquet/file_reader.h index fad5ee9eae38..9e2110f4a1e9 100644 --- a/cpp/src/parquet/file_reader.h +++ b/cpp/src/parquet/file_reader.h @@ -201,20 +201,13 @@ class PARQUET_EXPORT ParquetFileReader { const ::arrow::io::IOContext& ctx, const ::arrow::io::CacheOptions& options); - /// Release bytes buffered by a previous call to PreBuffer() for the - /// specified row groups and column indices. + /// \brief Release cached bytes pre-buffered for the given row groups/columns. /// - /// This releases the memory held by those cached ranges, allowing the - /// underlying buffers to be freed as soon as no other owner retains a - /// reference. Callers must ensure that the targeted row groups have been - /// fully decoded before invoking this method; otherwise subsequent reads - /// of the same row groups may fail because the cached bytes are gone. - /// - /// It is safe to call this concurrently with reads of other row groups: - /// cache entries that were coalesced across the targeted row groups and - /// unrelated data are left intact. - /// - /// Calling this method when PreBuffer() has not been called is a no-op. + /// Call only after the row groups are fully decoded; reads of those row + /// groups afterward may fail. An entry coalesced across row groups is freed + /// once all the row groups it spans have been evicted. Safe to call + /// concurrently for different row groups; a no-op if PreBuffer() was not + /// called. Returns the number of cache entries evicted. int64_t EvictPreBufferedData(const std::vector& row_groups, const std::vector& column_indices); From 8cecbeaebecbf283ecbf32ebe3244716896db7cf Mon Sep 17 00:00:00 2001 From: Justin Li Date: Tue, 16 Jun 2026 04:36:13 -0400 Subject: [PATCH 06/12] GH-39808: [C++][Parquet] Evict entries spanning filtered row groups --- .../parquet/arrow/arrow_reader_writer_test.cc | 45 ++++++++++ cpp/src/parquet/file_reader.cc | 83 +++++++++++++------ 2 files changed, 103 insertions(+), 25 deletions(-) diff --git a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc index 5071325ba3ab..4f6e3af18740 100644 --- a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc +++ b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc @@ -2869,6 +2869,51 @@ TEST(TestArrowReadWrite, EvictPreBufferedDataReleasesCrossRowGroupEntryOutOfOrde ASSERT_EQ(1, reader->parquet_reader()->EvictPreBufferedData({1}, column_indices)); } +// GH-39808: a filtered scan pre-buffers a non-contiguous subset of row groups. +// I/O coalescing can still merge column chunks of two buffered row groups into +// one entry that bridges the un-buffered (filtered-out) gap between them. Such +// an entry must be freed once both buffered row groups are evicted, even though +// they are not index-adjacent. +TEST(TestArrowReadWrite, EvictPreBufferedDataReleasesEntrySpanningFilteredRowGroups) { + ArrowReaderProperties properties = default_arrow_reader_properties(); + properties.set_pre_buffer(true); + const int num_rows = 1024; + const int row_group_size = 256; // 4 row groups + const int num_columns = 2; + const std::vector buffered_row_groups = {0, 2}; // 1 and 3 filtered out + const std::vector column_indices = {0, 1}; + + std::shared_ptr
table; + ASSERT_NO_FATAL_FAILURE(MakeDoubleTable(num_columns, num_rows, 1, &table)); + + std::shared_ptr buffer; + ASSERT_NO_FATAL_FAILURE(WriteTableToBuffer(table, row_group_size, + default_arrow_writer_properties(), &buffer)); + + std::unique_ptr reader; + FileReaderBuilder builder; + ASSERT_OK(builder.Open(std::make_shared(buffer))); + ASSERT_OK(builder.properties(properties)->Build(&reader)); + + // Huge limits coalesce RG0's and RG2's column chunks into ONE entry that + // bridges the un-buffered RG1 gap. + ::arrow::io::CacheOptions options = ::arrow::io::CacheOptions::LazyDefaults(); + options.hole_size_limit = static_cast(buffer->size()); + options.range_size_limit = static_cast(buffer->size()); + reader->parquet_reader()->PreBuffer(buffered_row_groups, column_indices, + ::arrow::io::IOContext(), options); + ASSERT_OK(reader->parquet_reader() + ->WhenBuffered(buffered_row_groups, column_indices) + .status()); + + // Evicting only RG0 leaves the entry in place (RG2 still needs it). + ASSERT_EQ(0, reader->parquet_reader()->EvictPreBufferedData({0}, column_indices)); + + // Evicting RG2 completes the buffered run {0, 2} and frees the spanning + // entry, even though the un-buffered RG1 lies between them by index. + ASSERT_EQ(1, reader->parquet_reader()->EvictPreBufferedData({2}, column_indices)); +} + TEST(TestArrowReadWrite, GetRecordBatchGenerator) { ArrowReaderProperties properties = default_arrow_reader_properties(); const int num_rows = 1024; diff --git a/cpp/src/parquet/file_reader.cc b/cpp/src/parquet/file_reader.cc index aa9bb6501049..d9a0f1618fe3 100644 --- a/cpp/src/parquet/file_reader.cc +++ b/cpp/src/parquet/file_reader.cc @@ -25,6 +25,7 @@ #include #include #include +#include #include #include #include @@ -384,10 +385,12 @@ class SerializedFile : public ParquetFileReader::Contents { cached_source_ = std::make_shared<::arrow::io::internal::ReadRangeCache>(source_, ctx, options); { - // GH-39808: a fresh cache invalidates any eviction runs we were - // tracking; they referenced the previous cache's byte ranges. + // GH-39808: a fresh cache invalidates any eviction runs we were tracking + // (they referenced the previous cache's byte ranges). Record the buffered + // row groups so eviction can merge runs across filtered-out gaps. std::lock_guard lock(eviction_mutex_); evicted_runs_.clear(); + buffered_row_groups_ = std::set(row_groups.begin(), row_groups.end()); } std::vector<::arrow::io::ReadRange> ranges; prebuffered_column_chunks_.clear(); @@ -442,11 +445,13 @@ class SerializedFile : public ParquetFileReader::Contents { } // Evict cached bytes populated by PreBuffer() for the given row groups. I/O - // coalescing can merge column chunks of adjacent row groups into one cache + // coalescing can merge column chunks of nearby row groups into one cache // entry, so an entry is freed only once every row group it covers has been - // evicted: we track contiguous runs of evicted row groups (row-group index - // order matches byte order) and evict each run's combined byte window. - // Returns the number of cache entries evicted. + // evicted: we track runs of evicted row groups that are contiguous in the + // buffered set and evict each run's combined byte window. A filtered-out + // (un-buffered) row group between two buffered ones does not block a merge; + // a buffered but not-yet-evicted neighbor does. Returns the number of cache + // entries evicted. // // Callers must have fully decoded the row groups (no reader still holds the // cached buffers). Safe to call concurrently for different row groups. @@ -478,23 +483,47 @@ class SerializedFile : public ParquetFileReader::Contents { continue; } - // Merge with the adjacent runs (one ending at row - 1, one starting at - // row + 1) into a single contiguous run, unioning the byte windows. + // Buffered row groups immediately before/after `row`. Runs are contiguous + // in the buffered set, so an un-buffered (filtered-out) gap is bridged + // while a buffered, not-yet-evicted neighbor blocks the merge. + bool have_pred = false; + bool have_succ = false; + int pred = 0; + int succ = 0; + auto row_it = buffered_row_groups_.find(row); + if (row_it != buffered_row_groups_.end()) { + if (row_it != buffered_row_groups_.begin()) { + pred = *std::prev(row_it); + have_pred = true; + } + auto succ_it = std::next(row_it); + if (succ_it != buffered_row_groups_.end()) { + succ = *succ_it; + have_succ = true; + } + } + + // Merge with the run starting at the buffered successor and/or the run + // ending at the buffered predecessor, unioning the byte windows. int first = row; int last = row; - auto right = evicted_runs_.find(row + 1); - if (right != evicted_runs_.end()) { - last = right->second.last_row_group; - end = std::max(end, right->second.end_offset); - evicted_runs_.erase(right); + if (have_succ) { + auto right = evicted_runs_.find(succ); + if (right != evicted_runs_.end()) { + last = right->second.last_row_group; + end = std::max(end, right->second.end_offset); + evicted_runs_.erase(right); + } } - auto upper = evicted_runs_.lower_bound(row); - if (upper != evicted_runs_.begin()) { - auto left = std::prev(upper); - if (left->second.last_row_group == row - 1) { - first = left->first; - start = std::min(start, left->second.start_offset); - evicted_runs_.erase(left); + if (have_pred) { + auto upper = evicted_runs_.lower_bound(row); + if (upper != evicted_runs_.begin()) { + auto left = std::prev(upper); + if (left->second.last_row_group == pred) { + first = left->first; + start = std::min(start, left->second.start_offset); + evicted_runs_.erase(left); + } } } evicted_runs_[first] = EvictedRowGroupRun{last, start, end}; @@ -688,11 +717,11 @@ class SerializedFile : public ParquetFileReader::Contents { std::shared_ptr page_index_reader_; std::unique_ptr bloom_filter_reader_; - // GH-39808: contiguous runs of row groups whose pre-buffered bytes have been - // evicted, keyed by each run's first row-group index. A coalesced cache entry - // can span adjacent row groups, so it is freed only once every row group in - // its run has been evicted. Guarded by eviction_mutex_ because row groups are - // decoded (and evicted) concurrently under readahead. + // GH-39808: runs of evicted row groups that are contiguous in the buffered + // set (so a run may span filtered-out row groups), keyed by each run's first + // row-group index. A coalesced cache entry is freed only once every row group + // in its run has been evicted. Guarded by eviction_mutex_ because row groups + // are decoded (and evicted) concurrently under readahead. struct EvictedRowGroupRun { int last_row_group; int64_t start_offset; @@ -700,6 +729,10 @@ class SerializedFile : public ParquetFileReader::Contents { }; std::mutex eviction_mutex_; std::map evicted_runs_; + // Row groups passed to the most recent PreBuffer(); lets eviction merge runs + // across filtered-out (un-buffered) row groups while still blocking on a + // buffered-but-undecoded neighbor. Guarded by eviction_mutex_. + std::set buffered_row_groups_; // Maps row group ordinal and prebuffer status of its column chunks in the form of a // bitmap buffer. std::unordered_map> prebuffered_column_chunks_; From 329eef60d9bc64f681bb1297c47d07f7628ea764 Mon Sep 17 00:00:00 2001 From: Justin Li Date: Sun, 28 Jun 2026 02:16:47 -0400 Subject: [PATCH 07/12] GH-39808: [C++] Replace EvictEntriesInRange with EvictEntriesBefore --- cpp/src/arrow/io/caching.cc | 27 +++----- cpp/src/arrow/io/caching.h | 21 +++---- cpp/src/arrow/io/memory_test.cc | 107 +++++++++++--------------------- 3 files changed, 54 insertions(+), 101 deletions(-) diff --git a/cpp/src/arrow/io/caching.cc b/cpp/src/arrow/io/caching.cc index ad13f192498d..3ad485af5339 100644 --- a/cpp/src/arrow/io/caching.cc +++ b/cpp/src/arrow/io/caching.cc @@ -265,29 +265,19 @@ struct ReadRangeCache::Impl { return AllComplete(futures); } - // Evict cache entries whose byte range is fully contained within - // [start_offset, start_offset + length). - Result EvictEntriesInRange(int64_t start_offset, int64_t length) { - if (length <= 0) { - return 0; - } - const int64_t end_offset = start_offset + length; + // Evict cache entries that end at or before `end_offset`. `entries` is sorted + // by offset, so we stop once an entry starts at/after `end_offset`. An entry + // that starts before but extends past `end_offset` (coalesced with a range a + // later consumer still needs) is left in place. + int64_t EvictEntriesBefore(int64_t end_offset) { int64_t n_evicted = 0; std::unique_lock guard(entry_mutex); - // entries is sorted by range.offset, so we can fast-forward to the first - // entry whose offset is >= start_offset and stop once we pass end_offset. - auto it = std::lower_bound(entries.begin(), entries.end(), start_offset, - [](const RangeCacheEntry& entry, int64_t offset) { - return entry.range.offset < offset; - }); + auto it = entries.begin(); while (it != entries.end() && it->range.offset < end_offset) { if (it->range.offset + it->range.length <= end_offset) { it = entries.erase(it); ++n_evicted; } else { - // Entry extends beyond the requested window (e.g. coalesced with a - // neighboring range). Leave it alone and keep scanning in case - // smaller siblings follow. ++it; } } @@ -378,9 +368,8 @@ Future<> ReadRangeCache::WaitFor(std::vector ranges) { return impl_->WaitFor(std::move(ranges)); } -Result ReadRangeCache::EvictEntriesInRange(int64_t start_offset, - int64_t length) { - return impl_->EvictEntriesInRange(start_offset, length); +int64_t ReadRangeCache::EvictEntriesBefore(int64_t end_offset) { + return impl_->EvictEntriesBefore(end_offset); } } // namespace internal diff --git a/cpp/src/arrow/io/caching.h b/cpp/src/arrow/io/caching.h index 97671d6b2b70..e3900944a560 100644 --- a/cpp/src/arrow/io/caching.h +++ b/cpp/src/arrow/io/caching.h @@ -142,20 +142,17 @@ class ARROW_EXPORT ReadRangeCache { /// \brief Wait until all given ranges have been cached. Future<> WaitFor(std::vector ranges); - /// \brief Evict cache entries whose byte range is fully contained within - /// [start_offset, start_offset + length). + /// \brief Evict cache entries that end at or before `end_offset`. /// - /// This releases the memory held by those entries, allowing buffers to be - /// freed as soon as no other owner retains a reference. Entries that are not - /// fully contained in the given window (for example, because I/O range - /// coalescing merged them with an adjacent region the caller still needs) - /// are not evicted, so this method is safe to call even if some cached - /// ranges have been merged across unrelated consumers. + /// Releases the memory of entries whose byte range lies entirely before + /// `end_offset`. An entry straddling `end_offset` is retained, so this is safe + /// even when I/O coalescing merged several requested ranges into one entry. + /// Buffers already returned by Read() stay valid through shared ownership. /// - /// \param[in] start_offset Start of the window (inclusive). - /// \param[in] length Length of the window. - /// \return Number of cache entries that were evicted. - Result EvictEntriesInRange(int64_t start_offset, int64_t length); + /// \param[in] end_offset Exclusive byte bound; entries ending at or before it + /// are evicted. + /// \return Number of cache entries evicted. + int64_t EvictEntriesBefore(int64_t end_offset); protected: struct Impl; diff --git a/cpp/src/arrow/io/memory_test.cc b/cpp/src/arrow/io/memory_test.cc index 812a4d133307..92a36e6a480a 100644 --- a/cpp/src/arrow/io/memory_test.cc +++ b/cpp/src/arrow/io/memory_test.cc @@ -919,78 +919,51 @@ TEST(RangeReadCache, LazyWithPrefetching) { ASSERT_RAISES(Invalid, cache.Read({25, 2})); } -TEST(RangeReadCache, EvictEntriesInRange) { +TEST(RangeReadCache, EvictEntriesBefore) { // GH-39808: entries cached by PreBuffer()-style code need to be evictable so - // that memory usage remains bounded while iterating over a large Parquet - // file. + // memory stays bounded while iterating a large Parquet file. std::string data = "abcdefghijklmnopqrstuvwxyz"; for (auto lazy : std::vector{false, true}) { SCOPED_TRACE(lazy); CacheOptions options = CacheOptions::Defaults(); - // Disable coalescing so each requested range becomes its own entry. - options.hole_size_limit = 0; + options.hole_size_limit = 0; // disable coalescing: one entry per range options.range_size_limit = 10; options.lazy = lazy; auto file = std::make_shared(std::make_shared(data)); internal::ReadRangeCache cache(file, {}, options); - // Cache three disjoint ranges, each forming its own entry. + // Entries: [1,3), [10,14), [20,22). ASSERT_OK(cache.Cache({{1, 2}, {10, 4}, {20, 2}})); - - // Sanity: all three ranges are readable. - ASSERT_OK_AND_ASSIGN(auto buf, cache.Read({1, 2})); - AssertBufferEqual(*buf, "bc"); - ASSERT_OK_AND_ASSIGN(buf, cache.Read({10, 4})); + ASSERT_OK_AND_ASSIGN(auto buf, cache.Read({10, 4})); AssertBufferEqual(*buf, "klmn"); - ASSERT_OK_AND_ASSIGN(buf, cache.Read({20, 2})); - AssertBufferEqual(*buf, "uv"); - // A window that doesn't fully contain any entry evicts nothing. - ASSERT_OK_AND_ASSIGN(int64_t evicted, cache.EvictEntriesInRange(0, 1)); - ASSERT_EQ(0, evicted); - ASSERT_OK_AND_ASSIGN(evicted, cache.EvictEntriesInRange(11, 2)); - ASSERT_EQ(0, evicted); - - // A zero- or negative-length window is a no-op. - ASSERT_OK_AND_ASSIGN(evicted, cache.EvictEntriesInRange(0, 0)); - ASSERT_EQ(0, evicted); - ASSERT_OK_AND_ASSIGN(evicted, cache.EvictEntriesInRange(10, -5)); - ASSERT_EQ(0, evicted); - - // Windows that partially overlap with an entry but don't fully contain - // it must also leave the entry in place. - ASSERT_OK_AND_ASSIGN(evicted, cache.EvictEntriesInRange(0, 2)); - ASSERT_EQ(0, evicted); - ASSERT_OK_AND_ASSIGN(evicted, cache.EvictEntriesInRange(10, 3)); - ASSERT_EQ(0, evicted); - - // Evicting the exact range of the middle entry removes only that entry. - ASSERT_OK_AND_ASSIGN(evicted, cache.EvictEntriesInRange(10, 4)); - ASSERT_EQ(1, evicted); - // Other entries are still readable. - ASSERT_OK_AND_ASSIGN(buf, cache.Read({1, 2})); - AssertBufferEqual(*buf, "bc"); - ASSERT_OK_AND_ASSIGN(buf, cache.Read({20, 2})); - AssertBufferEqual(*buf, "uv"); - // Evicted entry is gone. - ASSERT_RAISES(Invalid, cache.Read({10, 4})); + // An offset that splits no entry frees nothing. + ASSERT_EQ(0, cache.EvictEntriesBefore(0)); + ASSERT_EQ(0, cache.EvictEntriesBefore(2)); // [1,3) extends past 2 -> retained - // A wider window evicts every remaining fully-contained entry in one call. - ASSERT_OK_AND_ASSIGN(evicted, cache.EvictEntriesInRange(0, 100)); - ASSERT_EQ(2, evicted); + // end_offset == an entry's end frees exactly that entry ([1,3) ends at 3). + ASSERT_EQ(1, cache.EvictEntriesBefore(3)); ASSERT_RAISES(Invalid, cache.Read({1, 2})); + ASSERT_OK_AND_ASSIGN(buf, cache.Read({10, 4})); // others intact + AssertBufferEqual(*buf, "klmn"); + + // An offset inside an entry leaves it in place. + ASSERT_EQ(0, cache.EvictEntriesBefore(12)); // [10,14) straddles 12 + + // A wide offset frees every remaining entry in one call. + ASSERT_EQ(2, cache.EvictEntriesBefore(100)); + ASSERT_RAISES(Invalid, cache.Read({10, 4})); ASSERT_RAISES(Invalid, cache.Read({20, 2})); - // Evicting an already-empty cache is a safe no-op. - ASSERT_OK_AND_ASSIGN(evicted, cache.EvictEntriesInRange(0, 100)); - ASSERT_EQ(0, evicted); + // Empty cache is a safe no-op. + ASSERT_EQ(0, cache.EvictEntriesBefore(100)); } } TEST(RangeReadCache, ConcurrentReadAndEvict) { - // GH-39808: the Parquet dataset scanner calls EvictEntriesInRange from the + // GH-39808: the Parquet dataset scanner calls EvictEntriesBefore from the // thread-pool continuation that runs after a row group is decoded, while // other threads may still be calling Read() for column chunks of other // in-flight row groups. Exercise that pattern explicitly by slamming the @@ -1038,13 +1011,11 @@ TEST(RangeReadCache, ConcurrentReadAndEvict) { }; auto evictor_fn = [&]() { for (int iter = 0; iter < kIterations; ++iter) { - // Evict every lower-half entry. - for (int i = 0; i < kNumRanges / 2; ++i) { - auto result = cache.EvictEntriesInRange(ranges[i].offset, ranges[i].length); - if (!result.ok()) { - failures.fetch_add(1, std::memory_order_relaxed); - } - } + // Evict the entire lower half in one call (entries ending before the + // midpoint offset). + const int64_t mid_offset = + ranges[kNumRanges / 2 - 1].offset + ranges[kNumRanges / 2 - 1].length; + cache.EvictEntriesBefore(mid_offset); // Re-cache them so the next iteration has something to evict. std::vector lower(ranges.begin(), ranges.begin() + kNumRanges / 2); if (!cache.Cache(lower).ok()) { @@ -1068,33 +1039,29 @@ TEST(RangeReadCache, ConcurrentReadAndEvict) { } } -TEST(RangeReadCache, EvictEntriesInRangeSpanningEntry) { - // Coalesced entries (entries that were merged with a neighboring range at - // Cache() time) must not be evicted unless the eviction window fully - // contains them, otherwise we would drop bytes the caller still needs. +TEST(RangeReadCache, EvictEntriesBeforeSpanningEntry) { + // A coalesced entry must not be dropped until end_offset passes its end, + // otherwise we drop bytes a later consumer still needs. std::string data(40, 'x'); CacheOptions options = CacheOptions::Defaults(); - // Large hole_size_limit forces coalescing of adjacent ranges into one entry. - options.hole_size_limit = 100; + options.hole_size_limit = 100; // force coalescing into one entry options.range_size_limit = 200; auto file = std::make_shared(std::make_shared(data)); internal::ReadRangeCache cache(file, {}, options); - // These two ranges get coalesced into a single entry [1, 14). + // {1,3} and {10,4} coalesce into a single entry [1, 14). ASSERT_OK(cache.Cache({{1, 3}, {10, 4}})); - // Trying to evict only one of the two "logical" ranges must not drop the - // coalesced entry - it still holds data the other logical range needs. - ASSERT_OK_AND_ASSIGN(int64_t evicted, cache.EvictEntriesInRange(1, 3)); - ASSERT_EQ(0, evicted); + // An offset inside the entry (e.g. just past the first logical range) keeps it. + ASSERT_EQ(0, cache.EvictEntriesBefore(4)); + ASSERT_EQ(0, cache.EvictEntriesBefore(13)); ASSERT_OK_AND_ASSIGN(auto buf, cache.Read({10, 4})); ASSERT_EQ(4, buf->size()); - // A wide window that contains the whole coalesced entry evicts it. - ASSERT_OK_AND_ASSIGN(evicted, cache.EvictEntriesInRange(0, 20)); - ASSERT_EQ(1, evicted); + // end_offset at/after the entry's end (14) frees it. + ASSERT_EQ(1, cache.EvictEntriesBefore(14)); ASSERT_RAISES(Invalid, cache.Read({1, 3})); ASSERT_RAISES(Invalid, cache.Read({10, 4})); } From 81b96eef498c2c70d60492336e6546f8d098b875 Mon Sep 17 00:00:00 2001 From: Justin Li Date: Sun, 28 Jun 2026 02:35:47 -0400 Subject: [PATCH 08/12] GH-39808: [C++][Parquet] Evict pre-buffered bytes via offset watermark Replace the run-merging pre-buffer eviction with a simpler offset watermark. ParquetFileReader now exposes EvictPreBufferedDataBefore( end_offset), delegating to ReadRangeCache::EvictEntriesBefore, and the SerializedFile run-tracking state is removed. The async record-batch generator builds a ReadCacheEvictionState that advances a watermark over the contiguous prefix of decoded row groups and evicts every cache entry ending before the lowest byte any not-yet- completed row group still needs. A coalesced entry spanning a row-group boundary is freed once the watermark passes its end. Eviction is keyed off pre_buffer() and confined to this read-once generator path, so PreBuffer()'s contract for other callers is unchanged. --- .../parquet/arrow/arrow_reader_writer_test.cc | 207 +++++------------- cpp/src/parquet/arrow/reader.cc | 95 +++++++- cpp/src/parquet/file_reader.cc | 122 +---------- cpp/src/parquet/file_reader.h | 15 +- 4 files changed, 158 insertions(+), 281 deletions(-) diff --git a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc index 4f6e3af18740..24b87c221a66 100644 --- a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc +++ b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc @@ -24,8 +24,10 @@ #include "gmock/gmock.h" #include "gtest/gtest.h" +#include #include #include +#include #include #include #include @@ -2674,7 +2676,7 @@ TEST(TestArrowReadWrite, GetRecordBatchReaderNoColumns) { // should be released when the caller explicitly evicts them, otherwise // Dataset.to_batches accumulates memory across the lifetime of an open // FileReader. -TEST(TestArrowReadWrite, EvictPreBufferedData) { +TEST(TestArrowReadWrite, EvictPreBufferedDataBefore) { ArrowReaderProperties properties = default_arrow_reader_properties(); properties.set_pre_buffer(true); const int num_rows = 1024; @@ -2696,33 +2698,80 @@ TEST(TestArrowReadWrite, EvictPreBufferedData) { ASSERT_OK(builder.properties(properties)->Build(&reader)); ASSERT_EQ(reader->num_row_groups(), static_cast(row_groups.size())); - // Pre-buffer every row group/column so the cache is fully populated. reader->parquet_reader()->PreBuffer(row_groups, column_indices, ::arrow::io::IOContext(), ::arrow::io::CacheOptions::LazyDefaults()); ASSERT_OK(reader->parquet_reader()->WhenBuffered(row_groups, column_indices).status()); - // Decode the first row group; reads go through the cache. std::shared_ptr
rg_table; ASSERT_OK(reader->ReadRowGroup(/*i=*/0, column_indices, &rg_table)); ASSERT_EQ(rg_table->num_rows(), row_group_size); - // Evict only row group 0. The ranges for the other row groups are untouched, - // so they must still be readable. - reader->parquet_reader()->EvictPreBufferedData({0}, column_indices); + // The lowest byte offset row group 1 needs: evicting before it frees row + // group 0's (separately coalesced) entries while leaving row group 1 intact. + ASSERT_OK_AND_ASSIGN(auto rg1_ranges, + reader->parquet_reader()->GetReadRanges({1}, column_indices)); + int64_t rg1_min = std::numeric_limits::max(); + for (const auto& r : rg1_ranges) rg1_min = std::min(rg1_min, r.offset); + + EXPECT_GT(reader->parquet_reader()->EvictPreBufferedDataBefore(rg1_min), 0); ASSERT_OK(reader->ReadRowGroup(/*i=*/1, column_indices, &rg_table)); ASSERT_EQ(rg_table->num_rows(), row_group_size); - // Evicting twice is a no-op (and must not crash). - reader->parquet_reader()->EvictPreBufferedData({0}, column_indices); + // Re-evicting the same window frees nothing more. + ASSERT_EQ(0, reader->parquet_reader()->EvictPreBufferedDataBefore(rg1_min)); - // Calling EvictPreBufferedData on a reader that never called PreBuffer is - // also a no-op. + // A reader that never called PreBuffer is a no-op. std::unique_ptr no_prebuffer_reader; FileReaderBuilder no_prebuffer_builder; ASSERT_OK(no_prebuffer_builder.Open(std::make_shared(buffer))); ASSERT_OK(no_prebuffer_builder.Build(&no_prebuffer_reader)); - no_prebuffer_reader->parquet_reader()->EvictPreBufferedData(row_groups, column_indices); + ASSERT_EQ(0, + no_prebuffer_reader->parquet_reader()->EvictPreBufferedDataBefore(rg1_min)); +} + +// GH-39808: with coalescing a single cache entry can span adjacent row groups. +// Offset-based eviction frees such an entry once end_offset passes its end. +TEST(TestArrowReadWrite, EvictPreBufferedDataBeforeReleasesCrossRowGroupEntry) { + ArrowReaderProperties properties = default_arrow_reader_properties(); + properties.set_pre_buffer(true); + const int num_rows = 1024; + const int row_group_size = 256; // 4 row groups + const int num_columns = 2; + const std::vector row_groups = {0, 1, 2, 3}; + const std::vector column_indices = {0, 1}; + + std::shared_ptr
table; + ASSERT_NO_FATAL_FAILURE(MakeDoubleTable(num_columns, num_rows, 1, &table)); + + std::shared_ptr buffer; + ASSERT_NO_FATAL_FAILURE(WriteTableToBuffer(table, row_group_size, + default_arrow_writer_properties(), &buffer)); + + std::unique_ptr reader; + FileReaderBuilder builder; + ASSERT_OK(builder.Open(std::make_shared(buffer))); + ASSERT_OK(builder.properties(properties)->Build(&reader)); + + // Huge limits coalesce every column chunk of every row group into ONE entry + // spanning all row-group boundaries. + ::arrow::io::CacheOptions options = ::arrow::io::CacheOptions::LazyDefaults(); + options.hole_size_limit = static_cast(buffer->size()); + options.range_size_limit = static_cast(buffer->size()); + reader->parquet_reader()->PreBuffer(row_groups, column_indices, + ::arrow::io::IOContext(), options); + ASSERT_OK(reader->parquet_reader()->WhenBuffered(row_groups, column_indices).status()); + + // Any offset short of the spanning entry's end frees nothing. + ASSERT_OK_AND_ASSIGN(auto rg3_ranges, + reader->parquet_reader()->GetReadRanges({3}, column_indices)); + int64_t rg3_min = std::numeric_limits::max(); + for (const auto& r : rg3_ranges) rg3_min = std::min(rg3_min, r.offset); + ASSERT_EQ(0, reader->parquet_reader()->EvictPreBufferedDataBefore(rg3_min)); + + // An offset past the whole file frees the spanning entry. + ASSERT_EQ(1, reader->parquet_reader()->EvictPreBufferedDataBefore( + static_cast(buffer->size()))); } // GH-39808: when Dataset.to_batches-style iteration drives the async @@ -2778,142 +2827,6 @@ TEST(TestArrowReadWrite, GetRecordBatchGeneratorReleasesPreBufferedRowGroups) { AssertTablesEqual(*table, *actual, /*same_chunk_layout=*/false); } -// GH-39808: with default coalescing a single cache entry can span adjacent row -// groups. Per-row-group eviction must still release such an entry, but only -// once every row group it covers has been evicted. -TEST(TestArrowReadWrite, EvictPreBufferedDataReleasesCrossRowGroupEntry) { - ArrowReaderProperties properties = default_arrow_reader_properties(); - properties.set_pre_buffer(true); - const int num_rows = 1024; - const int row_group_size = 256; // 4 row groups - const int num_columns = 2; - const std::vector row_groups = {0, 1, 2, 3}; - const std::vector column_indices = {0, 1}; - - std::shared_ptr
table; - ASSERT_NO_FATAL_FAILURE(MakeDoubleTable(num_columns, num_rows, 1, &table)); - - std::shared_ptr buffer; - ASSERT_NO_FATAL_FAILURE(WriteTableToBuffer(table, row_group_size, - default_arrow_writer_properties(), &buffer)); - - std::unique_ptr reader; - FileReaderBuilder builder; - ASSERT_OK(builder.Open(std::make_shared(buffer))); - ASSERT_OK(builder.properties(properties)->Build(&reader)); - ASSERT_EQ(reader->num_row_groups(), static_cast(row_groups.size())); - - // Force every column chunk of every row group to coalesce into a single - // cache entry that spans all row-group boundaries. - ::arrow::io::CacheOptions options = ::arrow::io::CacheOptions::LazyDefaults(); - options.hole_size_limit = static_cast(buffer->size()); - options.range_size_limit = static_cast(buffer->size()); - reader->parquet_reader()->PreBuffer(row_groups, column_indices, - ::arrow::io::IOContext(), options); - ASSERT_OK(reader->parquet_reader()->WhenBuffered(row_groups, column_indices).status()); - - // Evicting any strict subset of the row groups leaves the spanning entry in - // place: it is not fully contained in any run that omits a row group it - // still covers, so each of these calls evicts nothing. - ASSERT_EQ(0, reader->parquet_reader()->EvictPreBufferedData({0}, column_indices)); - ASSERT_EQ(0, reader->parquet_reader()->EvictPreBufferedData({1}, column_indices)); - ASSERT_EQ(0, reader->parquet_reader()->EvictPreBufferedData({2}, column_indices)); - - // Once the final row group completes the contiguous run [0, 3], the spanning - // entry is fully contained and is freed. - ASSERT_EQ(1, reader->parquet_reader()->EvictPreBufferedData({3}, column_indices)); - - // Idempotent: re-evicting frees nothing more. - ASSERT_EQ(0, reader->parquet_reader()->EvictPreBufferedData({3}, column_indices)); -} - -// GH-39808: eviction order is non-deterministic under readahead (row groups -// decode concurrently). The spanning entry must be freed exactly once the gap -// between evicted runs is filled, regardless of order. -TEST(TestArrowReadWrite, EvictPreBufferedDataReleasesCrossRowGroupEntryOutOfOrder) { - ArrowReaderProperties properties = default_arrow_reader_properties(); - properties.set_pre_buffer(true); - const int num_rows = 1024; - const int row_group_size = 256; // 4 row groups - const int num_columns = 2; - const std::vector row_groups = {0, 1, 2, 3}; - const std::vector column_indices = {0, 1}; - - std::shared_ptr
table; - ASSERT_NO_FATAL_FAILURE(MakeDoubleTable(num_columns, num_rows, 1, &table)); - - std::shared_ptr buffer; - ASSERT_NO_FATAL_FAILURE(WriteTableToBuffer(table, row_group_size, - default_arrow_writer_properties(), &buffer)); - - std::unique_ptr reader; - FileReaderBuilder builder; - ASSERT_OK(builder.Open(std::make_shared(buffer))); - ASSERT_OK(builder.properties(properties)->Build(&reader)); - - ::arrow::io::CacheOptions options = ::arrow::io::CacheOptions::LazyDefaults(); - options.hole_size_limit = static_cast(buffer->size()); - options.range_size_limit = static_cast(buffer->size()); - reader->parquet_reader()->PreBuffer(row_groups, column_indices, - ::arrow::io::IOContext(), options); - ASSERT_OK(reader->parquet_reader()->WhenBuffered(row_groups, column_indices).status()); - - // Evict 2, then 0, then 3: runs {2,3} and {0} exist but neither covers the - // whole entry, so nothing is freed. - ASSERT_EQ(0, reader->parquet_reader()->EvictPreBufferedData({2}, column_indices)); - ASSERT_EQ(0, reader->parquet_reader()->EvictPreBufferedData({0}, column_indices)); - ASSERT_EQ(0, reader->parquet_reader()->EvictPreBufferedData({3}, column_indices)); - - // Evicting 1 fills the gap, merging {0} and {2,3} into {0,1,2,3}; the entry - // is now fully contained and freed. - ASSERT_EQ(1, reader->parquet_reader()->EvictPreBufferedData({1}, column_indices)); -} - -// GH-39808: a filtered scan pre-buffers a non-contiguous subset of row groups. -// I/O coalescing can still merge column chunks of two buffered row groups into -// one entry that bridges the un-buffered (filtered-out) gap between them. Such -// an entry must be freed once both buffered row groups are evicted, even though -// they are not index-adjacent. -TEST(TestArrowReadWrite, EvictPreBufferedDataReleasesEntrySpanningFilteredRowGroups) { - ArrowReaderProperties properties = default_arrow_reader_properties(); - properties.set_pre_buffer(true); - const int num_rows = 1024; - const int row_group_size = 256; // 4 row groups - const int num_columns = 2; - const std::vector buffered_row_groups = {0, 2}; // 1 and 3 filtered out - const std::vector column_indices = {0, 1}; - - std::shared_ptr
table; - ASSERT_NO_FATAL_FAILURE(MakeDoubleTable(num_columns, num_rows, 1, &table)); - - std::shared_ptr buffer; - ASSERT_NO_FATAL_FAILURE(WriteTableToBuffer(table, row_group_size, - default_arrow_writer_properties(), &buffer)); - - std::unique_ptr reader; - FileReaderBuilder builder; - ASSERT_OK(builder.Open(std::make_shared(buffer))); - ASSERT_OK(builder.properties(properties)->Build(&reader)); - - // Huge limits coalesce RG0's and RG2's column chunks into ONE entry that - // bridges the un-buffered RG1 gap. - ::arrow::io::CacheOptions options = ::arrow::io::CacheOptions::LazyDefaults(); - options.hole_size_limit = static_cast(buffer->size()); - options.range_size_limit = static_cast(buffer->size()); - reader->parquet_reader()->PreBuffer(buffered_row_groups, column_indices, - ::arrow::io::IOContext(), options); - ASSERT_OK(reader->parquet_reader() - ->WhenBuffered(buffered_row_groups, column_indices) - .status()); - - // Evicting only RG0 leaves the entry in place (RG2 still needs it). - ASSERT_EQ(0, reader->parquet_reader()->EvictPreBufferedData({0}, column_indices)); - - // Evicting RG2 completes the buffered run {0, 2} and frees the spanning - // entry, even though the un-buffered RG1 lies between them by index. - ASSERT_EQ(1, reader->parquet_reader()->EvictPreBufferedData({2}, column_indices)); -} - TEST(TestArrowReadWrite, GetRecordBatchGenerator) { ArrowReaderProperties properties = default_arrow_reader_properties(); const int num_rows = 1024; diff --git a/cpp/src/parquet/arrow/reader.cc b/cpp/src/parquet/arrow/reader.cc index 29a5a2065779..c0683577f507 100644 --- a/cpp/src/parquet/arrow/reader.cc +++ b/cpp/src/parquet/arrow/reader.cc @@ -19,7 +19,9 @@ #include #include +#include #include +#include #include #include #include @@ -1093,6 +1095,48 @@ Result> FileReaderImpl::GetRecordBatchReader( ::arrow::MakeFlattenIterator(std::move(batches)), std::move(batch_schema)); } +// GH-39808: releases pre-buffered cache entries as row groups are decoded. +// RowGroupDecoded() is called (possibly out of order, from readahead +// continuations) once a row group's batches are produced. It advances a +// watermark over the contiguous prefix of completed row groups and evicts every +// cache entry ending before the lowest byte any not-yet-completed row group +// still needs. A coalesced entry spanning a row-group boundary is freed once +// the watermark passes its end. +class ReadCacheEvictionState { + public: + // evict_before_offsets[i] = lowest byte offset row groups i..n-1 (in + // generator order) still need; evict_before_offsets[n] == INT64_MAX. + explicit ReadCacheEvictionState(std::vector evict_before_offsets) + : evict_before_offsets_(std::move(evict_before_offsets)), + completed_(evict_before_offsets_.size() - 1, false) {} + + void RowGroupDecoded(ParquetFileReader* reader, size_t row_group_index) { + int64_t evict_before = -1; + { + std::lock_guard lock(mutex_); + if (row_group_index >= completed_.size() || completed_[row_group_index]) { + return; + } + completed_[row_group_index] = true; + const size_t old_prefix = completed_prefix_; + while (completed_prefix_ < completed_.size() && completed_[completed_prefix_]) { + ++completed_prefix_; + } + if (completed_prefix_ == old_prefix) { + return; + } + evict_before = evict_before_offsets_[completed_prefix_]; + } + reader->EvictPreBufferedDataBefore(evict_before); + } + + private: + std::mutex mutex_; + std::vector evict_before_offsets_; + std::vector completed_; + size_t completed_prefix_ = 0; +}; + /// Given a file reader and a list of row groups, this is a generator of record /// batch generators (where each sub-generator is the contents of a single row group). class RowGroupGenerator { @@ -1108,12 +1152,14 @@ class RowGroupGenerator { explicit RowGroupGenerator(std::shared_ptr arrow_reader, ::arrow::internal::Executor* cpu_executor, std::vector row_groups, std::vector column_indices, - int64_t min_rows_in_flight) + int64_t min_rows_in_flight, + std::shared_ptr eviction_state) : arrow_reader_(std::move(arrow_reader)), cpu_executor_(cpu_executor), row_groups_(std::move(row_groups)), column_indices_(std::move(column_indices)), min_rows_in_flight_(min_rows_in_flight), + eviction_state_(std::move(eviction_state)), rows_in_flight_(0), index_(0), readahead_index_(0) {} @@ -1159,10 +1205,19 @@ class RowGroupGenerator { auto ready = reader->parquet_reader()->WhenBuffered({row_group}, column_indices); if (cpu_executor_) ready = cpu_executor_->TransferAlways(ready); row_group_read = - ready.Then([cpu_executor = cpu_executor_, reader, row_group, + ready.Then([cpu_executor = cpu_executor_, reader, row_group, row_group_index, + eviction_state = eviction_state_, column_indices = std::move( column_indices)]() -> ::arrow::Future { - return ReadOneRowGroup(cpu_executor, reader, row_group, column_indices); + return ReadOneRowGroup(cpu_executor, reader, row_group, column_indices) + .Then([reader, row_group_index, eviction_state]( + RecordBatchGenerator generator) -> RecordBatchGenerator { + if (eviction_state) { + eviction_state->RowGroupDecoded(reader->parquet_reader(), + row_group_index); + } + return generator; + }); }); } in_flight_reads_.push({std::move(row_group_read), num_rows}); @@ -1189,19 +1244,12 @@ class RowGroupGenerator { const int row_group, const std::vector& column_indices) { // Skips bound checks/pre-buffering, since we've done that already const int64_t batch_size = self->properties().batch_size(); - const bool pre_buffered = self->properties().pre_buffer(); return self->DecodeRowGroups(self, {row_group}, column_indices, cpu_executor) - .Then([batch_size, self, row_group, column_indices = column_indices, - pre_buffered](const std::shared_ptr
& table) + .Then([batch_size](const std::shared_ptr
& table) -> ::arrow::Result { ::arrow::TableBatchReader table_reader(*table); table_reader.set_chunksize(batch_size); ARROW_ASSIGN_OR_RAISE(auto batches, table_reader.ToRecordBatches()); - // GH-39808: release this row group's pre-buffered bytes now that it - // is decoded, keeping memory bounded while iterating. - if (pre_buffered) { - self->parquet_reader()->EvictPreBufferedData({row_group}, column_indices); - } return ::arrow::MakeVectorGenerator(std::move(batches)); }); } @@ -1211,6 +1259,7 @@ class RowGroupGenerator { std::vector row_groups_; std::vector column_indices_; int64_t min_rows_in_flight_; + std::shared_ptr eviction_state_; std::queue in_flight_reads_; int64_t rows_in_flight_; size_t index_; @@ -1233,10 +1282,32 @@ FileReaderImpl::GetRecordBatchGenerator(std::shared_ptr reader, reader_properties_.cache_options()); END_PARQUET_CATCH_EXCEPTIONS } + // GH-39808: when pre-buffering, release each row group's cached bytes once the + // contiguous prefix of decoded row groups no longer needs them, keeping the + // memory footprint bounded while iterating. Confined to this read-once path, + // so PreBuffer()'s contract for other callers is unchanged. + std::shared_ptr eviction_state; + if (reader_properties_.pre_buffer() && !column_indices.empty() && + !row_group_indices.empty()) { + const int64_t kNoMoreRanges = std::numeric_limits::max(); + std::vector evict_before(row_group_indices.size() + 1, kNoMoreRanges); + for (int64_t i = static_cast(row_group_indices.size()) - 1; i >= 0; --i) { + ARROW_ASSIGN_OR_RAISE( + auto ranges, reader_->GetReadRanges({row_group_indices[static_cast(i)]}, + column_indices)); + int64_t rg_min = kNoMoreRanges; + for (const auto& range : ranges) { + rg_min = std::min(rg_min, range.offset); + } + evict_before[static_cast(i)] = + std::min(evict_before[static_cast(i + 1)], rg_min); + } + eviction_state = std::make_shared(std::move(evict_before)); + } ::arrow::AsyncGenerator row_group_generator = RowGroupGenerator(::arrow::internal::checked_pointer_cast(reader), cpu_executor, row_group_indices, column_indices, - rows_to_readahead); + rows_to_readahead, std::move(eviction_state)); ::arrow::AsyncGenerator> concatenated = ::arrow::MakeConcatenatedGenerator(std::move(row_group_generator)); WRAP_ASYNC_GENERATOR(std::move(concatenated)); diff --git a/cpp/src/parquet/file_reader.cc b/cpp/src/parquet/file_reader.cc index d9a0f1618fe3..a3b6551124cd 100644 --- a/cpp/src/parquet/file_reader.cc +++ b/cpp/src/parquet/file_reader.cc @@ -384,14 +384,6 @@ class SerializedFile : public ParquetFileReader::Contents { const ::arrow::io::CacheOptions& options) { cached_source_ = std::make_shared<::arrow::io::internal::ReadRangeCache>(source_, ctx, options); - { - // GH-39808: a fresh cache invalidates any eviction runs we were tracking - // (they referenced the previous cache's byte ranges). Record the buffered - // row groups so eviction can merge runs across filtered-out gaps. - std::lock_guard lock(eviction_mutex_); - evicted_runs_.clear(); - buffered_row_groups_ = std::set(row_groups.begin(), row_groups.end()); - } std::vector<::arrow::io::ReadRange> ranges; prebuffered_column_chunks_.clear(); int num_cols = file_metadata_->num_columns(); @@ -444,95 +436,14 @@ class SerializedFile : public ParquetFileReader::Contents { return cached_source_->WaitFor(ranges); } - // Evict cached bytes populated by PreBuffer() for the given row groups. I/O - // coalescing can merge column chunks of nearby row groups into one cache - // entry, so an entry is freed only once every row group it covers has been - // evicted: we track runs of evicted row groups that are contiguous in the - // buffered set and evict each run's combined byte window. A filtered-out - // (un-buffered) row group between two buffered ones does not block a merge; - // a buffered but not-yet-evicted neighbor does. Returns the number of cache - // entries evicted. - // - // Callers must have fully decoded the row groups (no reader still holds the - // cached buffers). Safe to call concurrently for different row groups. - int64_t EvictPreBufferedData(const std::vector& row_groups, - const std::vector& column_indices) { - if (!cached_source_ || column_indices.empty()) { + // Evict cached bytes (from PreBuffer()) that end at or before `end_offset`. + // Buffers already handed to a decoder stay alive through shared ownership. + // A no-op if PreBuffer() was not called. + int64_t EvictPreBufferedDataBefore(int64_t end_offset) { + if (!cached_source_) { return 0; } - int64_t total_evicted = 0; - std::lock_guard lock(eviction_mutex_); - for (int row : row_groups) { - // Bounding box of this row group's selected column chunks. - int64_t start = std::numeric_limits::max(); - int64_t end = std::numeric_limits::min(); - for (int col : column_indices) { - auto range = - ComputeColumnChunkRange(file_metadata_.get(), source_size_, row, col); - start = std::min(start, range.offset); - end = std::max(end, range.offset + range.length); - } - if (end <= start) { - continue; - } - - // Skip row groups already covered by a run (idempotent). - auto after = evicted_runs_.upper_bound(row); - if (after != evicted_runs_.begin() && - row <= std::prev(after)->second.last_row_group) { - continue; - } - - // Buffered row groups immediately before/after `row`. Runs are contiguous - // in the buffered set, so an un-buffered (filtered-out) gap is bridged - // while a buffered, not-yet-evicted neighbor blocks the merge. - bool have_pred = false; - bool have_succ = false; - int pred = 0; - int succ = 0; - auto row_it = buffered_row_groups_.find(row); - if (row_it != buffered_row_groups_.end()) { - if (row_it != buffered_row_groups_.begin()) { - pred = *std::prev(row_it); - have_pred = true; - } - auto succ_it = std::next(row_it); - if (succ_it != buffered_row_groups_.end()) { - succ = *succ_it; - have_succ = true; - } - } - - // Merge with the run starting at the buffered successor and/or the run - // ending at the buffered predecessor, unioning the byte windows. - int first = row; - int last = row; - if (have_succ) { - auto right = evicted_runs_.find(succ); - if (right != evicted_runs_.end()) { - last = right->second.last_row_group; - end = std::max(end, right->second.end_offset); - evicted_runs_.erase(right); - } - } - if (have_pred) { - auto upper = evicted_runs_.lower_bound(row); - if (upper != evicted_runs_.begin()) { - auto left = std::prev(upper); - if (left->second.last_row_group == pred) { - first = left->first; - start = std::min(start, left->second.start_offset); - evicted_runs_.erase(left); - } - } - } - evicted_runs_[first] = EvictedRowGroupRun{last, start, end}; - - PARQUET_ASSIGN_OR_THROW(int64_t evicted, - cached_source_->EvictEntriesInRange(start, end - start)); - total_evicted += evicted; - } - return total_evicted; + return cached_source_->EvictEntriesBefore(end_offset); } // Metadata/footer parsing. Divided up to separate sync/async paths, and to use @@ -717,22 +628,6 @@ class SerializedFile : public ParquetFileReader::Contents { std::shared_ptr page_index_reader_; std::unique_ptr bloom_filter_reader_; - // GH-39808: runs of evicted row groups that are contiguous in the buffered - // set (so a run may span filtered-out row groups), keyed by each run's first - // row-group index. A coalesced cache entry is freed only once every row group - // in its run has been evicted. Guarded by eviction_mutex_ because row groups - // are decoded (and evicted) concurrently under readahead. - struct EvictedRowGroupRun { - int last_row_group; - int64_t start_offset; - int64_t end_offset; - }; - std::mutex eviction_mutex_; - std::map evicted_runs_; - // Row groups passed to the most recent PreBuffer(); lets eviction merge runs - // across filtered-out (un-buffered) row groups while still blocking on a - // buffered-but-undecoded neighbor. Guarded by eviction_mutex_. - std::set buffered_row_groups_; // Maps row group ordinal and prebuffer status of its column chunks in the form of a // bitmap buffer. std::unordered_map> prebuffered_column_chunks_; @@ -1025,12 +920,11 @@ void ParquetFileReader::PreBuffer(const std::vector& row_groups, file->PreBuffer(row_groups, column_indices, ctx, options); } -int64_t ParquetFileReader::EvictPreBufferedData(const std::vector& row_groups, - const std::vector& column_indices) { +int64_t ParquetFileReader::EvictPreBufferedDataBefore(int64_t end_offset) { // Access private methods here SerializedFile* file = ::arrow::internal::checked_cast(contents_.get()); - return file->EvictPreBufferedData(row_groups, column_indices); + return file->EvictPreBufferedDataBefore(end_offset); } Result> ParquetFileReader::GetReadRanges( diff --git a/cpp/src/parquet/file_reader.h b/cpp/src/parquet/file_reader.h index 9e2110f4a1e9..fa30d0fef814 100644 --- a/cpp/src/parquet/file_reader.h +++ b/cpp/src/parquet/file_reader.h @@ -201,15 +201,14 @@ class PARQUET_EXPORT ParquetFileReader { const ::arrow::io::IOContext& ctx, const ::arrow::io::CacheOptions& options); - /// \brief Release cached bytes pre-buffered for the given row groups/columns. + /// \brief Release cached bytes that end at or before `end_offset`. /// - /// Call only after the row groups are fully decoded; reads of those row - /// groups afterward may fail. An entry coalesced across row groups is freed - /// once all the row groups it spans have been evicted. Safe to call - /// concurrently for different row groups; a no-op if PreBuffer() was not - /// called. Returns the number of cache entries evicted. - int64_t EvictPreBufferedData(const std::vector& row_groups, - const std::vector& column_indices); + /// Only affects data cached by PreBuffer(). Call once the row groups needing + /// those bytes are fully decoded; later reads of evicted ranges may fail. + /// Buffers already handed to readers stay valid through shared ownership. A + /// no-op (returns 0) if PreBuffer() was not called. Returns the number of + /// cache entries evicted. + int64_t EvictPreBufferedDataBefore(int64_t end_offset); /// Retrieve the list of byte ranges that would need to be read to retrieve /// the data for the specified row groups and column indices. From a60afc81bb4e7bee846114f28747b5bc90abe53b Mon Sep 17 00:00:00 2001 From: Justin Li Date: Sun, 28 Jun 2026 11:34:42 -0400 Subject: [PATCH 09/12] GH-39808: [C++][Parquet] Test out-of-order pre-buffer eviction Move ReadCacheEvictionState to reader_internal.h and split its prefix-advance logic into a reader-free MarkDecodedAndGetEvictOffset() method that returns the evict-before offset (or nullopt). RowGroupDecoded is now a thin wrapper, so production behavior is unchanged while the out-of-order watermark advance is unit-testable without a reader. Add a deterministic ReadCacheEvictionStateOutOfOrder test driving that method directly, covering out-of-order completion, prefix jumps, and idempotent re-decode. Also drop now-dead includes ( in reader.cc once the class moved out; /// in file_reader.cc, left over from the removed run-merging members). --- .../parquet/arrow/arrow_reader_writer_test.cc | 18 ++++++ cpp/src/parquet/arrow/reader.cc | 43 -------------- cpp/src/parquet/arrow/reader_internal.h | 56 +++++++++++++++++++ cpp/src/parquet/file_reader.cc | 4 -- 4 files changed, 74 insertions(+), 47 deletions(-) diff --git a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc index 24b87c221a66..920970641911 100644 --- a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc +++ b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc @@ -28,6 +28,7 @@ #include #include #include +#include #include #include #include @@ -2827,6 +2828,23 @@ TEST(TestArrowReadWrite, GetRecordBatchGeneratorReleasesPreBufferedRowGroups) { AssertTablesEqual(*table, *actual, /*same_chunk_layout=*/false); } +// Readahead completes row groups out of order, so the watermark advance must +// only evict once the contiguous prefix of completed row groups grows. This +// drives that logic directly, without a reader, for determinism. +TEST(TestArrowReadWrite, ReadCacheEvictionStateOutOfOrder) { + constexpr int64_t kNoMoreRanges = std::numeric_limits::max(); + ReadCacheEvictionState state({100, 200, 300, kNoMoreRanges}); + + // Index 1 completes first: prefix stays at 0 (index 0 not done), no eviction. + EXPECT_EQ(state.MarkDecodedAndGetEvictOffset(1), std::nullopt); + // Index 0 completes: prefix jumps 0 -> 2, evict before row group 2's offset. + EXPECT_EQ(state.MarkDecodedAndGetEvictOffset(0), std::optional(300)); + // Index 2 completes: prefix jumps 2 -> 3, everything is evictable. + EXPECT_EQ(state.MarkDecodedAndGetEvictOffset(2), std::optional(kNoMoreRanges)); + // Re-decoding an already-completed index is idempotent: no eviction. + EXPECT_EQ(state.MarkDecodedAndGetEvictOffset(0), std::nullopt); +} + TEST(TestArrowReadWrite, GetRecordBatchGenerator) { ArrowReaderProperties properties = default_arrow_reader_properties(); const int num_rows = 1024; diff --git a/cpp/src/parquet/arrow/reader.cc b/cpp/src/parquet/arrow/reader.cc index c0683577f507..d7be8f07d8da 100644 --- a/cpp/src/parquet/arrow/reader.cc +++ b/cpp/src/parquet/arrow/reader.cc @@ -21,7 +21,6 @@ #include #include #include -#include #include #include #include @@ -1095,48 +1094,6 @@ Result> FileReaderImpl::GetRecordBatchReader( ::arrow::MakeFlattenIterator(std::move(batches)), std::move(batch_schema)); } -// GH-39808: releases pre-buffered cache entries as row groups are decoded. -// RowGroupDecoded() is called (possibly out of order, from readahead -// continuations) once a row group's batches are produced. It advances a -// watermark over the contiguous prefix of completed row groups and evicts every -// cache entry ending before the lowest byte any not-yet-completed row group -// still needs. A coalesced entry spanning a row-group boundary is freed once -// the watermark passes its end. -class ReadCacheEvictionState { - public: - // evict_before_offsets[i] = lowest byte offset row groups i..n-1 (in - // generator order) still need; evict_before_offsets[n] == INT64_MAX. - explicit ReadCacheEvictionState(std::vector evict_before_offsets) - : evict_before_offsets_(std::move(evict_before_offsets)), - completed_(evict_before_offsets_.size() - 1, false) {} - - void RowGroupDecoded(ParquetFileReader* reader, size_t row_group_index) { - int64_t evict_before = -1; - { - std::lock_guard lock(mutex_); - if (row_group_index >= completed_.size() || completed_[row_group_index]) { - return; - } - completed_[row_group_index] = true; - const size_t old_prefix = completed_prefix_; - while (completed_prefix_ < completed_.size() && completed_[completed_prefix_]) { - ++completed_prefix_; - } - if (completed_prefix_ == old_prefix) { - return; - } - evict_before = evict_before_offsets_[completed_prefix_]; - } - reader->EvictPreBufferedDataBefore(evict_before); - } - - private: - std::mutex mutex_; - std::vector evict_before_offsets_; - std::vector completed_; - size_t completed_prefix_ = 0; -}; - /// Given a file reader and a list of row groups, this is a generator of record /// batch generators (where each sub-generator is the contents of a single row group). class RowGroupGenerator { diff --git a/cpp/src/parquet/arrow/reader_internal.h b/cpp/src/parquet/arrow/reader_internal.h index 4ee3bf98bc54..c3e931b70308 100644 --- a/cpp/src/parquet/arrow/reader_internal.h +++ b/cpp/src/parquet/arrow/reader_internal.h @@ -22,6 +22,8 @@ #include #include #include +#include +#include #include #include #include @@ -133,5 +135,59 @@ Status TransferColumnData(::parquet::internal::RecordReader* reader, const ColumnDescriptor* descr, const ReaderContext* ctx, std::shared_ptr<::arrow::ChunkedArray>* out); +// ---------------------------------------------------------------------- +// Pre-buffer eviction + +// GH-39808: releases pre-buffered cache entries as row groups are decoded. +// RowGroupDecoded() is called (possibly out of order, from readahead +// continuations) once a row group's batches are produced. It advances a +// watermark over the contiguous prefix of completed row groups and evicts every +// cache entry ending before the lowest byte any not-yet-completed row group +// still needs. A coalesced entry spanning a row-group boundary is freed once +// the watermark passes its end. +class ReadCacheEvictionState { + public: + // evict_before_offsets[i] = lowest byte offset row groups i..n-1 (in + // generator order) still need; evict_before_offsets[n] == INT64_MAX. + // Must be non-empty: it always holds the n row-group offsets plus the + // INT64_MAX sentinel. + explicit ReadCacheEvictionState(std::vector evict_before_offsets) + : evict_before_offsets_(std::move(evict_before_offsets)), + completed_(evict_before_offsets_.size() - 1, false) {} + + void RowGroupDecoded(ParquetFileReader* reader, size_t row_group_index) { + if (auto evict_before = MarkDecodedAndGetEvictOffset(row_group_index)) { + reader->EvictPreBufferedDataBefore(*evict_before); + } + } + + // Marks row_group_index decoded and advances the contiguous completed prefix. + // Returns the offset to evict before when the prefix advanced, or std::nullopt + // when nothing newly became evictable (index out of range, already completed, + // or the prefix did not grow). Separated from RowGroupDecoded so the + // out-of-order watermark logic is unit-testable without a reader. + std::optional MarkDecodedAndGetEvictOffset(size_t row_group_index) { + std::lock_guard lock(mutex_); + if (row_group_index >= completed_.size() || completed_[row_group_index]) { + return std::nullopt; + } + completed_[row_group_index] = true; + const size_t old_prefix = completed_prefix_; + while (completed_prefix_ < completed_.size() && completed_[completed_prefix_]) { + ++completed_prefix_; + } + if (completed_prefix_ == old_prefix) { + return std::nullopt; + } + return evict_before_offsets_[completed_prefix_]; + } + + private: + std::mutex mutex_; + std::vector evict_before_offsets_; + std::vector completed_; + size_t completed_prefix_ = 0; +}; + } // namespace arrow } // namespace parquet diff --git a/cpp/src/parquet/file_reader.cc b/cpp/src/parquet/file_reader.cc index a3b6551124cd..1fbd0c7cd379 100644 --- a/cpp/src/parquet/file_reader.cc +++ b/cpp/src/parquet/file_reader.cc @@ -20,12 +20,8 @@ #include #include #include -#include -#include #include -#include #include -#include #include #include #include From de90ce731563257689a9dd6381f0671c79c1cc9c Mon Sep 17 00:00:00 2001 From: Justin Li Date: Sun, 28 Jun 2026 12:25:57 -0400 Subject: [PATCH 10/12] GH-39808: [C++][Parquet] Trim redundant comments Drop comments that restate the code or duplicate the header doc, keeping invariants and non-obvious rationale: the impl-side restatement above EvictEntriesBefore, the SerializedFile EvictPreBufferedDataBefore comment (already documented on the public method), and the non-empty precondition note on the eviction-state constructor (implied by the index invariant above it). --- cpp/src/arrow/io/caching.cc | 7 +++---- cpp/src/parquet/arrow/reader_internal.h | 2 -- cpp/src/parquet/file_reader.cc | 3 --- 3 files changed, 3 insertions(+), 9 deletions(-) diff --git a/cpp/src/arrow/io/caching.cc b/cpp/src/arrow/io/caching.cc index 3ad485af5339..af47173fe174 100644 --- a/cpp/src/arrow/io/caching.cc +++ b/cpp/src/arrow/io/caching.cc @@ -265,10 +265,9 @@ struct ReadRangeCache::Impl { return AllComplete(futures); } - // Evict cache entries that end at or before `end_offset`. `entries` is sorted - // by offset, so we stop once an entry starts at/after `end_offset`. An entry - // that starts before but extends past `end_offset` (coalesced with a range a - // later consumer still needs) is left in place. + // `entries` is sorted by offset, so we stop once an entry starts at/after + // `end_offset`. An entry that starts before but extends past `end_offset` + // (coalesced with a range a later consumer still needs) is left in place. int64_t EvictEntriesBefore(int64_t end_offset) { int64_t n_evicted = 0; std::unique_lock guard(entry_mutex); diff --git a/cpp/src/parquet/arrow/reader_internal.h b/cpp/src/parquet/arrow/reader_internal.h index c3e931b70308..c5bb18000480 100644 --- a/cpp/src/parquet/arrow/reader_internal.h +++ b/cpp/src/parquet/arrow/reader_internal.h @@ -149,8 +149,6 @@ class ReadCacheEvictionState { public: // evict_before_offsets[i] = lowest byte offset row groups i..n-1 (in // generator order) still need; evict_before_offsets[n] == INT64_MAX. - // Must be non-empty: it always holds the n row-group offsets plus the - // INT64_MAX sentinel. explicit ReadCacheEvictionState(std::vector evict_before_offsets) : evict_before_offsets_(std::move(evict_before_offsets)), completed_(evict_before_offsets_.size() - 1, false) {} diff --git a/cpp/src/parquet/file_reader.cc b/cpp/src/parquet/file_reader.cc index 1fbd0c7cd379..e0a412759a04 100644 --- a/cpp/src/parquet/file_reader.cc +++ b/cpp/src/parquet/file_reader.cc @@ -432,9 +432,6 @@ class SerializedFile : public ParquetFileReader::Contents { return cached_source_->WaitFor(ranges); } - // Evict cached bytes (from PreBuffer()) that end at or before `end_offset`. - // Buffers already handed to a decoder stay alive through shared ownership. - // A no-op if PreBuffer() was not called. int64_t EvictPreBufferedDataBefore(int64_t end_offset) { if (!cached_source_) { return 0; From 80da747fdf4cba6d069c80a7abac18b6c1139de0 Mon Sep 17 00:00:00 2001 From: Justin Li Date: Wed, 1 Jul 2026 09:03:05 +0900 Subject: [PATCH 11/12] GH-39808: [C++][Parquet] Fix cross-row-group test cache limits The cross-row-group eviction test set range_size_limit == hole_size_limit, which trips the range_size_limit > hole_size_limit check in debug/CI builds (the release build compiles it out). Make range one byte larger so the limits stay valid while still coalescing every chunk into one entry. Also trim the remaining added comments to bare invariants per review feedback. --- cpp/src/arrow/io/caching.cc | 5 ++--- cpp/src/arrow/io/caching.h | 14 ++++---------- .../parquet/arrow/arrow_reader_writer_test.cc | 8 +++----- cpp/src/parquet/arrow/reader.cc | 7 +++---- cpp/src/parquet/arrow/reader_internal.h | 17 +++++------------ cpp/src/parquet/file_reader.h | 10 +++------- 6 files changed, 20 insertions(+), 41 deletions(-) diff --git a/cpp/src/arrow/io/caching.cc b/cpp/src/arrow/io/caching.cc index af47173fe174..824afcad17fe 100644 --- a/cpp/src/arrow/io/caching.cc +++ b/cpp/src/arrow/io/caching.cc @@ -265,9 +265,8 @@ struct ReadRangeCache::Impl { return AllComplete(futures); } - // `entries` is sorted by offset, so we stop once an entry starts at/after - // `end_offset`. An entry that starts before but extends past `end_offset` - // (coalesced with a range a later consumer still needs) is left in place. + // `entries` is sorted by offset; an entry that extends past `end_offset` + // (coalesced with a range a later consumer still needs) is kept. int64_t EvictEntriesBefore(int64_t end_offset) { int64_t n_evicted = 0; std::unique_lock guard(entry_mutex); diff --git a/cpp/src/arrow/io/caching.h b/cpp/src/arrow/io/caching.h index e3900944a560..700ed919228d 100644 --- a/cpp/src/arrow/io/caching.h +++ b/cpp/src/arrow/io/caching.h @@ -142,16 +142,10 @@ class ARROW_EXPORT ReadRangeCache { /// \brief Wait until all given ranges have been cached. Future<> WaitFor(std::vector ranges); - /// \brief Evict cache entries that end at or before `end_offset`. - /// - /// Releases the memory of entries whose byte range lies entirely before - /// `end_offset`. An entry straddling `end_offset` is retained, so this is safe - /// even when I/O coalescing merged several requested ranges into one entry. - /// Buffers already returned by Read() stay valid through shared ownership. - /// - /// \param[in] end_offset Exclusive byte bound; entries ending at or before it - /// are evicted. - /// \return Number of cache entries evicted. + /// \brief Evict cache entries ending at or before `end_offset`, returning the + /// count. An entry straddling `end_offset` is retained (safe when coalescing + /// merged ranges). Buffers already returned by Read() stay valid through + /// shared ownership. int64_t EvictEntriesBefore(int64_t end_offset); protected: diff --git a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc index 920970641911..0b10bef9af78 100644 --- a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc +++ b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc @@ -2673,10 +2673,8 @@ TEST(TestArrowReadWrite, GetRecordBatchReaderNoColumns) { ASSERT_EQ(actual_batch->num_rows(), num_rows); } -// GH-39808: bytes cached by PreBuffer() for an already-decoded row group -// should be released when the caller explicitly evicts them, otherwise -// Dataset.to_batches accumulates memory across the lifetime of an open -// FileReader. +// GH-39808: bytes cached by PreBuffer() for a decoded row group must be +// releasable, else Dataset.to_batches accumulates memory over the reader's life. TEST(TestArrowReadWrite, EvictPreBufferedDataBefore) { ArrowReaderProperties properties = default_arrow_reader_properties(); properties.set_pre_buffer(true); @@ -2758,7 +2756,7 @@ TEST(TestArrowReadWrite, EvictPreBufferedDataBeforeReleasesCrossRowGroupEntry) { // spanning all row-group boundaries. ::arrow::io::CacheOptions options = ::arrow::io::CacheOptions::LazyDefaults(); options.hole_size_limit = static_cast(buffer->size()); - options.range_size_limit = static_cast(buffer->size()); + options.range_size_limit = static_cast(buffer->size()) + 1; reader->parquet_reader()->PreBuffer(row_groups, column_indices, ::arrow::io::IOContext(), options); ASSERT_OK(reader->parquet_reader()->WhenBuffered(row_groups, column_indices).status()); diff --git a/cpp/src/parquet/arrow/reader.cc b/cpp/src/parquet/arrow/reader.cc index d7be8f07d8da..1092f27a06c1 100644 --- a/cpp/src/parquet/arrow/reader.cc +++ b/cpp/src/parquet/arrow/reader.cc @@ -1239,10 +1239,9 @@ FileReaderImpl::GetRecordBatchGenerator(std::shared_ptr reader, reader_properties_.cache_options()); END_PARQUET_CATCH_EXCEPTIONS } - // GH-39808: when pre-buffering, release each row group's cached bytes once the - // contiguous prefix of decoded row groups no longer needs them, keeping the - // memory footprint bounded while iterating. Confined to this read-once path, - // so PreBuffer()'s contract for other callers is unchanged. + // GH-39808: evict each row group's bytes as the decoded prefix advances, so + // memory stays bounded. Only this read-once path evicts, so PreBuffer()'s + // contract is unchanged for other callers. std::shared_ptr eviction_state; if (reader_properties_.pre_buffer() && !column_indices.empty() && !row_group_indices.empty()) { diff --git a/cpp/src/parquet/arrow/reader_internal.h b/cpp/src/parquet/arrow/reader_internal.h index c5bb18000480..d36c376d8eaa 100644 --- a/cpp/src/parquet/arrow/reader_internal.h +++ b/cpp/src/parquet/arrow/reader_internal.h @@ -138,13 +138,9 @@ Status TransferColumnData(::parquet::internal::RecordReader* reader, // ---------------------------------------------------------------------- // Pre-buffer eviction -// GH-39808: releases pre-buffered cache entries as row groups are decoded. -// RowGroupDecoded() is called (possibly out of order, from readahead -// continuations) once a row group's batches are produced. It advances a -// watermark over the contiguous prefix of completed row groups and evicts every -// cache entry ending before the lowest byte any not-yet-completed row group -// still needs. A coalesced entry spanning a row-group boundary is freed once -// the watermark passes its end. +// GH-39808: as row groups finish decoding (possibly out of order under +// readahead), advance a watermark over the leading run of completed ones and +// evict cache entries ending before the lowest byte any remaining one needs. class ReadCacheEvictionState { public: // evict_before_offsets[i] = lowest byte offset row groups i..n-1 (in @@ -159,11 +155,8 @@ class ReadCacheEvictionState { } } - // Marks row_group_index decoded and advances the contiguous completed prefix. - // Returns the offset to evict before when the prefix advanced, or std::nullopt - // when nothing newly became evictable (index out of range, already completed, - // or the prefix did not grow). Separated from RowGroupDecoded so the - // out-of-order watermark logic is unit-testable without a reader. + // Reader-free half of RowGroupDecoded (so the out-of-order advance is testable): + // the evict-before offset when the prefix grows, else nullopt. std::optional MarkDecodedAndGetEvictOffset(size_t row_group_index) { std::lock_guard lock(mutex_); if (row_group_index >= completed_.size() || completed_[row_group_index]) { diff --git a/cpp/src/parquet/file_reader.h b/cpp/src/parquet/file_reader.h index fa30d0fef814..1c7407fe5fa8 100644 --- a/cpp/src/parquet/file_reader.h +++ b/cpp/src/parquet/file_reader.h @@ -201,13 +201,9 @@ class PARQUET_EXPORT ParquetFileReader { const ::arrow::io::IOContext& ctx, const ::arrow::io::CacheOptions& options); - /// \brief Release cached bytes that end at or before `end_offset`. - /// - /// Only affects data cached by PreBuffer(). Call once the row groups needing - /// those bytes are fully decoded; later reads of evicted ranges may fail. - /// Buffers already handed to readers stay valid through shared ownership. A - /// no-op (returns 0) if PreBuffer() was not called. Returns the number of - /// cache entries evicted. + /// \brief Release cached bytes (from PreBuffer()) ending at or before + /// `end_offset`. Call once those row groups are decoded; later reads of evicted + /// ranges may fail. No-op (returns 0) if PreBuffer() was not called. int64_t EvictPreBufferedDataBefore(int64_t end_offset); /// Retrieve the list of byte ranges that would need to be read to retrieve From 522f1dffa8079808290734d2b466eaec1300c720 Mon Sep 17 00:00:00 2001 From: Justin Li Date: Thu, 16 Jul 2026 04:04:03 -0400 Subject: [PATCH 12/12] GH-39808: [C++][Parquet] Fall back after read-cache eviction --- cpp/src/arrow/io/caching.cc | 17 +++++++---------- .../parquet/arrow/arrow_reader_writer_test.cc | 7 +++++++ cpp/src/parquet/file_reader.cc | 12 +++++++++--- cpp/src/parquet/file_reader.h | 3 ++- 4 files changed, 25 insertions(+), 14 deletions(-) diff --git a/cpp/src/arrow/io/caching.cc b/cpp/src/arrow/io/caching.cc index 824afcad17fe..5d7df22b738b 100644 --- a/cpp/src/arrow/io/caching.cc +++ b/cpp/src/arrow/io/caching.cc @@ -268,17 +268,14 @@ struct ReadRangeCache::Impl { // `entries` is sorted by offset; an entry that extends past `end_offset` // (coalesced with a range a later consumer still needs) is kept. int64_t EvictEntriesBefore(int64_t end_offset) { - int64_t n_evicted = 0; std::unique_lock guard(entry_mutex); - auto it = entries.begin(); - while (it != entries.end() && it->range.offset < end_offset) { - if (it->range.offset + it->range.length <= end_offset) { - it = entries.erase(it); - ++n_evicted; - } else { - ++it; - } - } + const auto first_kept = + std::remove_if(entries.begin(), entries.end(), + [end_offset](const RangeCacheEntry& entry) { + return entry.range.offset + entry.range.length <= end_offset; + }); + const auto n_evicted = static_cast(entries.end() - first_kept); + entries.erase(first_kept, entries.end()); return n_evicted; } diff --git a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc index 0b10bef9af78..9e3e67edf38b 100644 --- a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc +++ b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc @@ -2717,6 +2717,13 @@ TEST(TestArrowReadWrite, EvictPreBufferedDataBefore) { ASSERT_OK(reader->ReadRowGroup(/*i=*/1, column_indices, &rg_table)); ASSERT_EQ(rg_table->num_rows(), row_group_size); + EXPECT_GT(reader->parquet_reader()->EvictPreBufferedDataBefore( + std::numeric_limits::max()), + 0); + std::shared_ptr reread_column; + ASSERT_OK(reader->RowGroup(1)->Column(0)->Read(&reread_column)); + ASSERT_EQ(reread_column->length(), row_group_size); + // Re-evicting the same window frees nothing more. ASSERT_EQ(0, reader->parquet_reader()->EvictPreBufferedDataBefore(rg1_min)); diff --git a/cpp/src/parquet/file_reader.cc b/cpp/src/parquet/file_reader.cc index e0a412759a04..0711b479d794 100644 --- a/cpp/src/parquet/file_reader.cc +++ b/cpp/src/parquet/file_reader.cc @@ -247,9 +247,15 @@ class SerializedRowGroup : public RowGroupReader::Contents { ::arrow::bit_util::GetBit(prebuffered_column_chunks_bitmap_->data(), i)) { // PARQUET-1698: if read coalescing is enabled, read from pre-buffered // segments. - PARQUET_ASSIGN_OR_THROW(auto buffer, cached_source_->Read(col_range)); - stream = std::make_shared<::arrow::io::BufferReader>(buffer); - } else { + auto buffer = cached_source_->Read(col_range); + if (!buffer.ok() && !buffer.status().IsInvalid()) { + PARQUET_THROW_NOT_OK(buffer.status()); + } + if (buffer.ok()) { + stream = std::make_shared<::arrow::io::BufferReader>(*std::move(buffer)); + } + } + if (stream == nullptr) { stream = properties_.GetStream(source_, col_range.offset, col_range.length); } diff --git a/cpp/src/parquet/file_reader.h b/cpp/src/parquet/file_reader.h index 1c7407fe5fa8..8bae39e6127f 100644 --- a/cpp/src/parquet/file_reader.h +++ b/cpp/src/parquet/file_reader.h @@ -203,7 +203,8 @@ class PARQUET_EXPORT ParquetFileReader { /// \brief Release cached bytes (from PreBuffer()) ending at or before /// `end_offset`. Call once those row groups are decoded; later reads of evicted - /// ranges may fail. No-op (returns 0) if PreBuffer() was not called. + /// ranges fall back to the source file. No-op (returns 0) if PreBuffer() was not + /// called. int64_t EvictPreBufferedDataBefore(int64_t end_offset); /// Retrieve the list of byte ranges that would need to be read to retrieve