[justinli500]

Rationale for this change

Dataset Parquet scans can retain prebuffered row-group data until the file reader is destroyed. This increases peak memory use during batch iteration.

I added eviction for cached read ranges to address this part of #39808. Other sources of Dataset memory growth are outside this PR.

What changes are included in this PR?

• I added ReadRangeCache::EvictEntriesBefore() to release entries ending at or before a safe byte offset. Entries spanning that offset stay cached, and buffers already returned to readers remain valid.
• I advance that offset in row-group output order, keeping the bytes needed by remaining row groups. This handles any row-group order; reading in file order releases memory sooner.
• I enabled automatic eviction for Dataset Parquet scans. I left it disabled by default for general FileReader users, who may read the same row groups again.
• I changed ReadRangeCache::Read() to return an optional buffer and updated its callers. A cache miss returns an empty optional, while I/O errors still propagate. Parquet falls back to the source on a miss, and IPC keeps its existing Invalid error for a miss.
• I kept vector storage and used an index to track live entries, so eviction can clear expired futures without shifting surviving entries. I also synchronized cache reads, waits, prefetching, and eviction.

Are these changes tested?

I ran the affected C++ suites locally in Release:

Test suite	Result
arrow-io-memory-test	59 passed
arrow-ipc-read-write-test	410 passed
arrow-dataset-file-parquet-test	60 passed
parquet-arrow-reader-writer-test	833 passed, 2 skipped, 1 disabled

I also ran the cache suite and focused Dataset and Parquet eviction checks in a Debug build with threading disabled. The concurrent cache test skips in that configuration.

I added a Dataset scan regression that checks cached compressed bytes are released during iteration and verifies the output. I kept the lower-level tests for coalesced entries, rereads, default cache retention, and consumed and upcoming ranges. I checked that the regressions fail when eviction is disabled or delayed until the last row group.

I syntax-checked the changed S3 benchmark read function, but I have not built the full S3 benchmark target. I have not rerun end-to-end memory or throughput benchmarks for this revision.

Are there any user-facing changes?

I added ParquetFileReader::EvictPreBufferedDataBefore() and the ArrowReaderProperties::auto_evict_read_cache property to the C++ API. Dataset Parquet scans enable eviction internally; general FileReader defaults stay the same.

I changed the return type of the internal ReadRangeCache::Read() API to an optional buffer.

I used AI tools to help with the implementation, tests, and review replies.

• GitHub Issue: [Python] Dataset.to_batches accumulates memory usage and leaks #39808

[github-actions]

⚠️ GitHub issue #39808 has been automatically assigned in GitHub to PR creator.

[justinli500]

Just pushed a clang-format fix for the only PR-related CI failure (lint). The two Windows MinGW failures should be an unrelated arrow-s3fs-test flake. Could a reviewer approve CI and take a look? Thanks!

[bveeramani]

@wgtmac @pitrou This is a huge issue for Ray users, and we're trying to decide if we need to move away from PyArrow for our Parquet reading implementation

Will you guys be unable to prioritize this PR? I know maintaining a project is a lot of work, so I'd understand if you can't take a look immediately

[wgtmac]

Thanks for the fix! I have two concerns with this patch.

First, eviction does not seem to account for coalesced cache entries that span row groups. See my inline comment.

Second, many of the comments feel too verbose and restate what the code is already doing. I’d suggest trimming them down to only the invariants or non-obvious reasoning.

[wgtmac]

Eviction is done one row group at a time, but cache entries are only removed if they are fully contained in that row group’s byte window. With default coalescing, a single cache entry can span adjacent row groups, so evicting row group 0 leaves the entry because it extends past the window, and evicting row group 1 also leaves it because the entry starts before the window. That entry is then never released, so the memory growth this PR is meant to fix can still occur for small or adjacent row groups.

[justinli500]

Thanks for the feedback! Both addressed:

• Cross-RG eviction: tracks runs of evicted row groups and evicts each run's combined window, so that a coalesced entry is freed once all its row groups are evicted (EvictEntriesInRange unchanged). Merges by buffered-set adjacency, so filtered/non-contiguous scans are covered too; safe under out-of-order eviction. Tests added for all three cases.
• Comments: trimmed to the invariants.

[wgtmac]

Thanks for updating it! I still think this PR is more complex than necessary. A few concerns:

• The new eviction API requires callers to manually decide when and what to evict, which is easy to misuse.
• The implementation tracks evicted state by row group runs, but the API also accepts column_indices, so column-subset eviction can be incorrect or ambiguous.
• The run-merging / bounding-box logic is subtle, especially with filtered row groups and coalesced cache entries.
• The comments are quite verbose and sometimes explain the implementation too mechanically. It makes the patch feel harder to review than the actual idea.

Could we instead make this automatic for the sequential Arrow reader path? For example, add an opt-in reader property, and after each row group is decoded, advance a simple watermark over the contiguous completed row groups. Then evict cache entries ending before the first byte any remaining row group may need.

That would keep PreBuffer() semantics unchanged, avoid a public row-group/column eviction API, handle readahead safely, and reduce the amount of Parquet-specific bookkeeping needed.

I asked Codex to draft the patch below to demonstrate my idea. I'm not suggesting to adopt it as is but just a demo code in case I didn't make it clear.

Details

diff --git a/cpp/src/arrow/io/caching.cc b/cpp/src/arrow/io/caching.cc
index 41fdd9f781..8305869bcf 100644
--- a/cpp/src/arrow/io/caching.cc
+++ b/cpp/src/arrow/io/caching.cc
@@ -151,8 +151,10 @@ 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).
+  // Protects entries and the futures stored in entries.
   std::vector<RangeCacheEntry> entries;
+  std::mutex mutex;
 
   virtual ~Impl() = default;
 
@@ -179,14 +181,17 @@ struct ReadRangeCache::Impl {
         ranges, internal::CoalesceReadRanges(std::move(ranges), options.hole_size_limit,
                                              options.range_size_limit));
     std::vector<RangeCacheEntry> new_entries = MakeCacheEntries(ranges);
-    // Add new entries, themselves ordered by offset
-    if (entries.size() > 0) {
-      std::vector<RangeCacheEntry> 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);
+    {
+      std::unique_lock<std::mutex> guard(mutex);
+      // Add new entries, themselves ordered by offset
+      if (entries.size() > 0) {
+        std::vector<RangeCacheEntry> 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);
@@ -205,14 +210,20 @@ struct ReadRangeCache::Impl {
       return std::make_shared<Buffer>(&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<std::shared_ptr<Buffer>> fut;
+    int64_t slice_offset = 0;
+    {
+      std::unique_lock<std::mutex> guard(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;
@@ -226,19 +237,41 @@ 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");
+    ARROW_ASSIGN_OR_RAISE(auto buf, fut.result());
+    return SliceBuffer(std::move(buf), slice_offset, range.length);
   }
 
   virtual Future<> Wait() {
     std::vector<Future<>> futures;
-    for (auto& entry : entries) {
-      futures.emplace_back(MaybeRead(&entry));
+    {
+      std::unique_lock<std::mutex> guard(mutex);
+      futures.reserve(entries.size());
+      for (auto& entry : entries) {
+        futures.emplace_back(MaybeRead(&entry));
+      }
     }
     return AllComplete(futures);
   }
 
+  int64_t EvictEntriesBefore(int64_t end_offset) {
+    int64_t n_evicted = 0;
+    std::unique_lock<std::mutex> guard(mutex);
+    auto it = entries.begin();
+    while (it != entries.end()) {
+      if (it->range.offset >= end_offset) {
+        break;
+      }
+      if (it->range.length <= end_offset - it->range.offset) {
+        it = entries.erase(it);
+        ++n_evicted;
+      } else {
+        ++it;
+      }
+    }
+    return n_evicted;
+  }
+
   // Return a Future that completes when the given ranges have been read.
   virtual Future<> WaitFor(std::vector<ReadRange> ranges) {
     auto end = std::remove_if(ranges.begin(), ranges.end(),
@@ -246,17 +279,21 @@ struct ReadRangeCache::Impl {
     ranges.resize(end - ranges.begin());
     std::vector<Future<>> 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<std::mutex> guard(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);
@@ -339,6 +376,10 @@ Future<> ReadRangeCache::WaitFor(std::vector<ReadRange> ranges) {
   return impl_->WaitFor(std::move(ranges));
 }
 
+int64_t ReadRangeCache::EvictEntriesBefore(int64_t end_offset) {
+  return impl_->EvictEntriesBefore(end_offset);
+}
+
 }  // namespace internal
 }  // namespace io
 }  // namespace arrow
diff --git a/cpp/src/arrow/io/caching.h b/cpp/src/arrow/io/caching.h
index e2b911fafd..55b1a03fc1 100644
--- a/cpp/src/arrow/io/caching.h
+++ b/cpp/src/arrow/io/caching.h
@@ -142,6 +142,14 @@ class ARROW_EXPORT ReadRangeCache {
   /// \brief Wait until all given ranges have been cached.
   Future<> WaitFor(std::vector<ReadRange> ranges);
 
+  /// \brief Evict cache entries ending at or before the given offset.
+  ///
+  /// This releases the memory held by those entries. Buffers already returned
+  /// by Read() remain valid through their shared ownership.
+  ///
+  /// \return Number of cache entries that were evicted.
+  int64_t EvictEntriesBefore(int64_t end_offset);
+
  protected:
   struct Impl;
   struct LazyImpl;
diff --git a/cpp/src/parquet/arrow/reader.cc b/cpp/src/parquet/arrow/reader.cc
index cc107c1802..9e1fbbb538 100644
--- a/cpp/src/parquet/arrow/reader.cc
+++ b/cpp/src/parquet/arrow/reader.cc
@@ -19,7 +19,9 @@
 
 #include <algorithm>
 #include <cstring>
+#include <limits>
 #include <memory>
+#include <mutex>
 #include <random>
 #include <unordered_set>
 #include <utility>
@@ -1144,6 +1146,39 @@ Result<std::unique_ptr<RecordBatchReader>> FileReaderImpl::GetRecordBatchReader(
 
 /// 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 ReadCacheEvictionState {
+ public:
+  explicit ReadCacheEvictionState(std::vector<int64_t> 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<std::mutex> lock(mutex_);
+      if (row_group_index >= completed_.size() || completed_[row_group_index]) {
+        return;
+      }
+      completed_[row_group_index] = true;
+      const size_t old_completed_prefix = completed_prefix_;
+      while (completed_prefix_ < completed_.size() && completed_[completed_prefix_]) {
+        ++completed_prefix_;
+      }
+      if (completed_prefix_ == old_completed_prefix) {
+        return;
+      }
+      evict_before = evict_before_offsets_[completed_prefix_];
+    }
+    reader->EvictPreBufferedDataBefore(evict_before);
+  }
+
+ private:
+  std::mutex mutex_;
+  std::vector<int64_t> evict_before_offsets_;
+  std::vector<bool> completed_;
+  size_t completed_prefix_ = 0;
+};
+
 class RowGroupGenerator {
  public:
   using RecordBatchGenerator =
@@ -1157,12 +1192,14 @@ class RowGroupGenerator {
   explicit RowGroupGenerator(std::shared_ptr<FileReaderImpl> arrow_reader,
                              ::arrow::internal::Executor* cpu_executor,
                              std::vector<int> row_groups, std::vector<int> column_indices,
-                             int64_t min_rows_in_flight)
+                             int64_t min_rows_in_flight,
+                             std::shared_ptr<ReadCacheEvictionState> 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) {}
@@ -1207,12 +1244,21 @@ class RowGroupGenerator {
     } else {
       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,
-                      column_indices = std::move(
-                          column_indices)]() -> ::arrow::Future<RecordBatchGenerator> {
-            return ReadOneRowGroup(cpu_executor, reader, row_group, column_indices);
-          });
+      row_group_read = ready.Then([cpu_executor = cpu_executor_, reader, row_group,
+                                   row_group_index, eviction_state = eviction_state_,
+                                   column_indices = std::move(column_indices)]()
+                                      -> ::arrow::Future<RecordBatchGenerator> {
+        return ReadOneRowGroup(cpu_executor, reader, row_group, column_indices)
+            .Then([reader, row_group_index, eviction_state = std::move(eviction_state)](
+                      RecordBatchGenerator generator)
+                      -> ::arrow::Result<RecordBatchGenerator> {
+              if (eviction_state) {
+                eviction_state->RowGroupDecoded(reader->parquet_reader(),
+                                                row_group_index);
+              }
+              return std::move(generator);
+            });
+      });
     }
     in_flight_reads_.push({std::move(row_group_read), num_rows});
   }
@@ -1253,6 +1299,7 @@ class RowGroupGenerator {
   std::vector<int> row_groups_;
   std::vector<int> column_indices_;
   int64_t min_rows_in_flight_;
+  std::shared_ptr<ReadCacheEvictionState> eviction_state_;
   std::queue<ReadRequest> in_flight_reads_;
   int64_t rows_in_flight_;
   size_t index_;
@@ -1275,10 +1322,30 @@ FileReaderImpl::GetRecordBatchGenerator(std::shared_ptr<FileReader> reader,
                        reader_properties_.cache_options());
     END_PARQUET_CATCH_EXCEPTIONS
   }
+  std::shared_ptr<ReadCacheEvictionState> eviction_state;
+  if (reader_properties_.pre_buffer() && reader_properties_.auto_evict_read_cache() &&
+      !column_indices.empty()) {
+    const int64_t no_more_ranges = std::numeric_limits<int64_t>::max();
+    std::vector<int64_t> evict_before_offsets(row_group_indices.size() + 1,
+                                              no_more_ranges);
+    for (int64_t i = static_cast<int64_t>(row_group_indices.size()) - 1; i >= 0; --i) {
+      ARROW_ASSIGN_OR_RAISE(
+          auto ranges, reader_->GetReadRanges({row_group_indices[static_cast<size_t>(i)]},
+                                              column_indices));
+      int64_t row_group_min_offset = no_more_ranges;
+      for (const auto& range : ranges) {
+        row_group_min_offset = std::min(row_group_min_offset, range.offset);
+      }
+      evict_before_offsets[static_cast<size_t>(i)] = std::min(
+          evict_before_offsets[static_cast<size_t>(i + 1)], row_group_min_offset);
+    }
+    eviction_state =
+        std::make_shared<ReadCacheEvictionState>(std::move(evict_before_offsets));
+  }
   ::arrow::AsyncGenerator<RowGroupGenerator::RecordBatchGenerator> row_group_generator =
       RowGroupGenerator(::arrow::internal::checked_pointer_cast<FileReaderImpl>(reader),
                         cpu_executor, row_group_indices, column_indices,
-                        rows_to_readahead);
+                        rows_to_readahead, std::move(eviction_state));
   ::arrow::AsyncGenerator<std::shared_ptr<::arrow::RecordBatch>> 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 2f46a5e296..bd6c3f3115 100644
--- a/cpp/src/parquet/file_reader.cc
+++ b/cpp/src/parquet/file_reader.cc
@@ -403,6 +403,13 @@ class SerializedFile : public ParquetFileReader::Contents {
     PARQUET_THROW_NOT_OK(cached_source_->Cache(ranges));
   }
 
+  int64_t EvictPreBufferedDataBefore(int64_t end_offset) {
+    if (!cached_source_) {
+      return 0;
+    }
+    return cached_source_->EvictEntriesBefore(end_offset);
+  }
+
   Result<std::vector<::arrow::io::ReadRange>> GetReadRanges(
       const std::vector<int>& row_groups, const std::vector<int>& column_indices,
       int64_t hole_size_limit, int64_t range_size_limit) {
@@ -909,6 +916,13 @@ void ParquetFileReader::PreBuffer(const std::vector<int>& row_groups,
   file->PreBuffer(row_groups, column_indices, ctx, options);
 }
 
+int64_t ParquetFileReader::EvictPreBufferedDataBefore(int64_t end_offset) {
+  // Access private methods here
+  SerializedFile* file =
+      ::arrow::internal::checked_cast<SerializedFile*>(contents_.get());
+  return file->EvictPreBufferedDataBefore(end_offset);
+}
+
 Result<std::vector<::arrow::io::ReadRange>> ParquetFileReader::GetReadRanges(
     const std::vector<int>& row_groups, const std::vector<int>& 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 c42163276c..8662dcc0d9 100644
--- a/cpp/src/parquet/file_reader.h
+++ b/cpp/src/parquet/file_reader.h
@@ -201,6 +201,12 @@ class PARQUET_EXPORT ParquetFileReader {
                  const ::arrow::io::IOContext& ctx,
                  const ::arrow::io::CacheOptions& options);
 
+  /// \brief Release cached entries ending at or before the given offset.
+  ///
+  /// This only affects data cached by PreBuffer(). Buffers already returned to
+  /// readers remain valid through their shared ownership.
+  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.
   ///
diff --git a/cpp/src/parquet/properties.h b/cpp/src/parquet/properties.h
index e2244a1176..abd22683fb 100644
--- a/cpp/src/parquet/properties.h
+++ b/cpp/src/parquet/properties.h
@@ -1151,6 +1151,7 @@ class PARQUET_EXPORT ArrowReaderProperties {
         read_dict_indices_(),
         batch_size_(kArrowDefaultBatchSize),
         pre_buffer_(true),
+        auto_evict_read_cache_(false),
         cache_options_(::arrow::io::CacheOptions::LazyDefaults()),
         coerce_int96_timestamp_unit_(::arrow::TimeUnit::NANO),
         binary_type_(kArrowDefaultBinaryType),
@@ -1235,6 +1236,14 @@ class PARQUET_EXPORT ArrowReaderProperties {
   /// Return whether read coalescing is enabled.
   bool pre_buffer() const { return pre_buffer_; }
 
+  /// Set whether to automatically evict pre-buffered read cache entries while
+  /// consuming row groups sequentially.
+  ///
+  /// Default is false.
+  void set_auto_evict_read_cache(bool auto_evict) { auto_evict_read_cache_ = auto_evict; }
+  /// Return whether the Arrow reader automatically evicts read cache entries.
+  bool auto_evict_read_cache() const { return auto_evict_read_cache_; }
+
   /// Set options for read coalescing. This can be used to tune the
   /// implementation for characteristics of different filesystems.
   void set_cache_options(::arrow::io::CacheOptions options) { cache_options_ = options; }
@@ -1300,6 +1309,7 @@ class PARQUET_EXPORT ArrowReaderProperties {
   std::unordered_set<int> read_dict_indices_;
   int64_t batch_size_;
   bool pre_buffer_;
+  bool auto_evict_read_cache_;
   ::arrow::io::IOContext io_context_;
   ::arrow::io::CacheOptions cache_options_;
   ::arrow::TimeUnit::type coerce_int96_timestamp_unit_;

[wgtmac]

@pitrou Do you have an opinion on this? I don't want to waste author's effort if my suggestion is not the right direction.

[wgtmac]

cc @mapleFU

[justinli500]

Received - @wgtmac your approach looks great and I'd be glad to pivot.

Just wondering, though, since it defaults off, how would Dataset.to_batches in the #39808 path actually get the fix? Would we enable it in the scanner, or flip the default later?

Happy to hold until we get more eyes on this though.

[wgtmac]

I made the option default off only to avoid changing the general PreBuffer() contract. Some callers may expect the cache to stay alive until the next PreBuffer() or reader destruction. For the Dataset GetRecordBatchGenerator path, I think we should enable this automatically when pre_buffer=true. That path reads row groups once, so it has a clear safe point to evict after each row group is decoded.

[justinli500]

Got it - in that case, I'll pivot to your approach and start working on it now. Going to key eviction off pre_buffer() in the GetRecordBatchGenerator path so the Dataset/#39808 case is covered, and will use the offset-watermark so that merged cross-row-group entries free by construction.

[wgtmac]

Could you please make CI happy and trim those added comments (they are still too long to read)?

[justinli500]

Just fixed both, should work now!

[rok]

It's great to see this being worked on!
I'm not very familiar with this part of the codebase so some more general comments.

[rok]

@wgtmac would you like to do another pass?

[wgtmac]

I‘ll take another round this week. Thanks for the reminder! @rok

[wgtmac]

Suggested change

	// -- Polymorphic hooks. Always called with entry_mutex held. --

These comments look more like notes while doing the implementation.

[justinli500]

Agreed, removed.

[wgtmac]

Suggested change

	// 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::deque<RangeCacheEntry> entries;
	std::deque<RangeCacheEntry> entries; // GUARDED_BY(entry_mutex)

No strong opinion. I would prefer simplifying comments like this.

[justinli500]

Sounds good, just shortened this section.

[wgtmac]

Suggested change

// -- Public entry points (acquire entry_mutex, then delegate). --

[wgtmac]

Why do we need to return a int64_t? It would be good to document its meaning or just make it void if not very useful.

[justinli500]

Good point - I changed this to void and updated the tests.

[wgtmac]

This looks hacky to fallback on a specific error code and may hide its original error. Is it better to introduce a function like Result<std::optional<std::shared_ptr<Buffer>>> ReadIfCached(ReadRange range) to explicitly indicate that cache miss is not an error?

[justinli500]

Good point - I added ReadIfCached so we can tell a cache miss from a real error. We only fall back on a miss now, and I also added a test for the error case.

[wgtmac]

Suggested change

	if (entries.size() > 0) {
	std::deque<RangeCacheEntry> 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 {
	for (auto& entry : new_entries) {
	entries.push_back(std::move(entry));
	}
	}
	std::deque<RangeCacheEntry> merged;
	std::merge(std::make_move_iterator(entries.begin()),
	std::make_move_iterator(entries.end()),
	std::make_move_iterator(new_entries.begin()),
	std::make_move_iterator(new_entries.end()),
	std::back_inserter(merged));
	entries.swap(merged);

This might be more efficient.

[HuaHuaY]

Why not to use entries.insert(entries.begin(), std::make_move_iterator(new_entries.begin()), std::make_move_iterator(new_entries.end()))?

[justinli500]

Great point! I went with the move-based merge version - the new ranges can fall between existing ones, so inserting them all at the front could definitely cause some issues.

[wgtmac]

Would it be more useful to return total evicted memory?

[justinli500]

I made this void instead. Nothing uses the return value, and a byte count could be misleading if something else still holds the buffers.

[wgtmac]

As I've said earlier, we cannot evict them by default. There are use cases where users create a file reader by prebuffering some row groups and repeatedly call GetRecordBatchGenerator to read the same row group for more than 1 time. In that case, followup reads may have issues. We can only enable it by default to the Dataset parquet reader where each row group is read only once.

[justinli500]

Ah right sorry about that - I changed it so this is only enabled by the Dataset reader, and normal file readers keep the existing behavior. I added a test for that too.

[wgtmac]

I think we can remove ReadCacheEvictionState if eviction is performed in row-group output order instead of decode-completion order.
GetRecordBatchGenerator() can keep the existing suffix-min calculation, but move the resulting vector directly into RowGroupGenerator:

RowGroupGenerator(..., std::vector<int64_t> evict_before_offsets)
    : ...,
      evict_before_offsets_(std::move(evict_before_offsets)) {}

std::vector<int64_t> evict_before_offsets_;

FetchNext() would only create and enqueue row_group_read. Eviction can instead be attached when operator() pops the corresponding FIFO request:

Future<RecordBatchGenerator> operator()() {
  if (index_ >= row_groups_.size()) {
    return AsyncGeneratorEnd<RecordBatchGenerator>();
  }

  const size_t request_index = index_++;
  FillReadahead();

  DCHECK(!in_flight_reads_.empty());
  ReadRequest next = std::move(in_flight_reads_.front());
  in_flight_reads_.pop();
  rows_in_flight_ -= next.num_rows;

  if (evict_before_offsets_.empty()) {
    return next.read;
  }

  auto reader = arrow_reader_;
  return next.read.Then(
      [reader, offset = evict_before_offsets_[request_index + 1]](
          RecordBatchGenerator generator) {
        reader->parquet_reader()->EvictPreBufferedDataBefore(offset);
        return generator;
      });
}

The queue already preserves row-group order, and MakeConcatenatedGenerator consumes these generators sequentially. A read completed early by readahead would therefore retain its compressed buffer until its output turn, but retention remains bounded by the readahead window.

[justinli500]

Thanks, I went with this. Eviction now follows the queued output order, and I also removed ReadCacheEvictionState.

[HuaHuaY]

Does MakeCacheEntries need to be protected by entry_mutex?

[justinli500]

Good point, it doesn’t. I moved it outside the lock and kept only the shared cache update inside.

[HuaHuaY]

Suggested change

	int64_t n_evicted = 0;
	while (!entries.empty()) {
	const auto& range = entries.front().range;
	if (range.offset + range.length > end_offset) {
	break;
	}
	entries.pop_front();
	++n_evicted;
	}
	const auto first_kept =
	std::find_if(entries.cbegin(), entries.cend(), [end_offset](const RangeCacheEntry& entry) {
	return entry.range.offset + entry.range.length > end_offset;
	});
	const auto n_evicted = static_cast<int64_t>(first_kept - entries.begin());
	entries.erase(entries.begin(), first_kept);

You can determine the quantity first, and then delete them all at once.

[justinli500]

Thanks, I changed it to find the cutoff once and erase the prefix in one go. I left out the count since the method now returns void.

[jwilliams-ocient]

Thanks for all of your efforts to maintain this library!

@rok , @wgtmac, @HuaHuaY, is there a target release for this improvement?

[wgtmac]

@jwilliams-ocient Sorry for the long delay! I will take a look again.

[wgtmac]

@justinli500 Sorry for the long delay! I haven't carefully read new test case yet. I think the general approach looks good to me with some question on minor issues.

I believe this issue is worth fixing. Do you want to take a quick look at this? @adamreeve @pitrou

[wgtmac]

ReadRangeCache is under internal namespace. Perhaps we can directly change the Read signature above to make it return Result<std::optional<std::shared_ptr<Buffer>>> in order not to add one more function?

WDYT? @pitrou

[justinli500]

Sounds good. I changed Read() to return an optional buffer and removed ReadIfCached(). I updated the callers, and IPC still reports a cache miss as Invalid.

[wgtmac]

I'm not sure if changing vector to deque will have a performance regression for the default non-evicting readers.

If we still use vector, eviction could clear the expired futures to drop the cache’s buffer references, then advance the index without shifting surviving entries. Lookups, waits, prefetching, and merges would operate on the live suffix, with compaction deferred or batched.

[justinli500]

Good point. I switched back to vector and used an index to track the live entries. Now eviction clears expired futures, and Cache() merges the remaining entries with new ranges.

[wgtmac]

nit: add a DCHECK for evict_before_offsets_.size() just in case?

[wgtmac]

This change seems unrelated?

[wgtmac]

I think we need at least one test case for the dataset reader.

[justinli500]

Agreed, I added a Dataset scan test for this. It checks that cached bytes are released during iteration and that the output is correct. It also fails if Dataset doesn’t enable eviction.

[wgtmac]

ReadRowGroup() calls PreBuffer() again, replacing the multi-row-group cache with just this row group. Please use a read path that preserves the cache and assert that the upcoming row group remains cached after eviction.

[justinli500]

Good catch - I disabled automatic prebuffering in this test so each read preserves the cache created by the explicit PreBuffer() call. Now it checks that upcoming groups remain cached after eviction.

[wgtmac]

Codex has popped up this review comment to me: please guard this test with ARROW_ENABLE_THREADING, like the existing threading tests. It currently starts five std::threads even in threadless configurations.

[wgtmac]

This also passes if eviction is delayed until the last row group. Please force separate cache entries and check consumed and upcoming ranges during iteration.

[justinli500]

I updated this to use separate cache entries and check both consumed and upcoming ranges during iteration. I also checked that the test fails if eviction is delayed until the last group.

[wgtmac]

I'm not sure if it is better to use a real dataset reader test case to replace this one.

[justinli500]

I added the Dataset test and kept the FileReader test as well. The FileReader test still covers default cache retention and explicit eviction.

[adamreeve]

This looks like a useful change to me thanks @justinli500.

I think this could also be re-used to fix #46935 as a follow-up.

[adamreeve]

I think this comment could be improved to make it easier to understand. How about this:

This is intended for one-pass consumers, and reduces memory use by freeing buffered ranges
that are no longer needed. It is safe to enable when reading row groups in any order, but reading
them in file order allows the cache to release memory progressively. Default is false.