GH-39808: [C++][Parquet] Evict pre-buffered row-group bytes after decode

cpp/src/arrow/dataset/file_parquet.cc

@@ -130,6 +130,8 @@ parquet::ArrowReaderProperties MakeArrowReaderProperties(
   // Must be set here since the sync ScanTask handles pre-buffering itself
   arrow_properties.set_pre_buffer(
       parquet_scan_options.arrow_reader_properties->pre_buffer());
+  // Dataset scans consume each row group once, so cached ranges can be released.
+  arrow_properties.set_auto_evict_read_cache(true);
   arrow_properties.set_cache_options(
       parquet_scan_options.arrow_reader_properties->cache_options());
   arrow_properties.set_io_context(

cpp/src/arrow/io/caching.cc

@@ -18,7 +18,10 @@
 #include <algorithm>
 #include <atomic>
 #include <cmath>
+#include <deque>
+#include <iterator>
 #include <mutex>
+#include <optional>
 #include <utility>
 #include <vector>
 
@@ -151,12 +154,12 @@ struct ReadRangeCache::Impl {
   IOContext ctx;
   CacheOptions options;
 
-  // Ordered by offset (so as to find a matching region by binary search)
-  std::vector<RangeCacheEntry> entries;
+  std::deque<RangeCacheEntry> entries;  // GUARDED_BY(entry_mutex)
+  std::mutex entry_mutex;
 
   virtual ~Impl() = default;
 
-  // Get the future corresponding to a range
+  // Get the future corresponding to a range. Called with entry_mutex held.
   virtual Future<std::shared_ptr<Buffer>> MaybeRead(RangeCacheEntry* entry) {
     return entry->future;
   }
@@ -173,46 +176,55 @@ struct ReadRangeCache::Impl {
     return new_entries;
   }
 
-  // Add the given ranges to the cache, coalescing them where possible
-  virtual Status Cache(std::vector<ReadRange> ranges) {
+  // Add the given ranges to the cache, coalescing them where possible.
+  Status Cache(std::vector<ReadRange> ranges) {
     ARROW_ASSIGN_OR_RAISE(
         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);
+    Status st;
+    {
+      std::unique_lock<std::mutex> guard(entry_mutex);
+      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);
     }
-    // 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();
     }
     return st;
   }
 
-  // Read the given range from the cache, blocking if needed. Cannot read a range
-  // that spans cache entries.
-  virtual Result<std::shared_ptr<Buffer>> Read(ReadRange range) {
+  // Read the given range from the cache, blocking if needed. Cannot read a
+  // range that spans cache entries.
+  Result<std::optional<std::shared_ptr<Buffer>>> ReadIfCached(ReadRange range) {
     if (range.length == 0) {
       static const uint8_t byte = 0;
-      return std::make_shared<Buffer>(&byte, 0);
+      return std::make_optional(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(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 std::nullopt;
+      }
+      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,53 +238,78 @@ 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 std::make_optional(SliceBuffer(std::move(buf), slice_offset, range.length));
+  }
+
+  Result<std::shared_ptr<Buffer>> Read(ReadRange range) {
+    ARROW_ASSIGN_OR_RAISE(auto buffer, ReadIfCached(range));
+    if (!buffer) {
+      return Status::Invalid("ReadRangeCache did not find matching cache entry");
+    }
+    return std::move(*buffer);
   }
 
-  virtual Future<> Wait() {
+  Future<> Wait() {
     std::vector<Future<>> futures;
-    for (auto& entry : entries) {
-      futures.emplace_back(MaybeRead(&entry));
+    {
+      std::unique_lock<std::mutex> guard(entry_mutex);
+      futures.reserve(entries.size());
+      for (auto& entry : entries) {
+        futures.emplace_back(MaybeRead(&entry));
+      }
     }
     return AllComplete(futures);
   }
 
+  // Cached ranges are sorted and non-overlapping, so entries ending at or
+  // before `end_offset` form a prefix. Keep a straddling entry.
+  void EvictEntriesBefore(int64_t end_offset) {
+    std::unique_lock<std::mutex> guard(entry_mutex);
+    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;
+        });
+    entries.erase(entries.cbegin(), first_kept);
+  }
+
   // Return a Future that completes when the given ranges have been read.
-  virtual Future<> WaitFor(std::vector<ReadRange> ranges) {
+  Future<> WaitFor(std::vector<ReadRange> ranges) {
     auto end = std::remove_if(ranges.begin(), ranges.end(),
                               [](const ReadRange& range) { return range.length == 0; });
     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(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).
 struct ReadRangeCache::LazyImpl : public ReadRangeCache::Impl {
-  // Protect against concurrent modification of entries[i]->future
-  std::mutex entry_mutex;
-
   virtual ~LazyImpl() = default;
 
   Future<std::shared_ptr<Buffer>> MaybeRead(RangeCacheEntry* entry) override {
-    // Called by superclass Read()/WaitFor() so we have the lock
     if (!entry->future.is_valid()) {
       entry->future = file->ReadAsync(ctx, entry->range.offset, entry->range.length,
                                       /*allow_short_read=*/false);
@@ -291,26 +328,6 @@ struct ReadRangeCache::LazyImpl : public ReadRangeCache::Impl {
     }
     return new_entries;
   }
-
-  Status Cache(std::vector<ReadRange> ranges) override {
-    std::unique_lock<std::mutex> guard(entry_mutex);
-    return ReadRangeCache::Impl::Cache(std::move(ranges));
-  }
-
-  Result<std::shared_ptr<Buffer>> Read(ReadRange range) override {
-    std::unique_lock<std::mutex> guard(entry_mutex);
-    return ReadRangeCache::Impl::Read(range);
-  }
-
-  Future<> Wait() override {
-    std::unique_lock<std::mutex> guard(entry_mutex);
-    return ReadRangeCache::Impl::Wait();
-  }
-
-  Future<> WaitFor(std::vector<ReadRange> ranges) override {
-    std::unique_lock<std::mutex> guard(entry_mutex);
-    return ReadRangeCache::Impl::WaitFor(std::move(ranges));
-  }
 };
 
 ReadRangeCache::ReadRangeCache(std::shared_ptr<RandomAccessFile> owned_file,
@@ -333,12 +350,21 @@ Result<std::shared_ptr<Buffer>> ReadRangeCache::Read(ReadRange range) {
   return impl_->Read(range);
 }
 
+Result<std::optional<std::shared_ptr<Buffer>>> ReadRangeCache::ReadIfCached(
+    ReadRange range) {
+  return impl_->ReadIfCached(range);
+}
+
 Future<> ReadRangeCache::Wait() { return impl_->Wait(); }
 
 Future<> ReadRangeCache::WaitFor(std::vector<ReadRange> ranges) {
   return impl_->WaitFor(std::move(ranges));
 }
 
+void ReadRangeCache::EvictEntriesBefore(int64_t end_offset) {
+  impl_->EvictEntriesBefore(end_offset);
+}
+
 }  // namespace internal
 }  // namespace io
 }  // namespace arrow

cpp/src/arrow/io/caching.h

@@ -19,6 +19,7 @@
 
 #include <cstdint>
 #include <memory>
+#include <optional>
 #include <string>
 #include <utility>
 #include <vector>
@@ -136,12 +137,23 @@ class ARROW_EXPORT ReadRangeCache {
   /// \brief Read a range previously given to Cache().
   Result<std::shared_ptr<Buffer>> Read(ReadRange range);
 
+  /// \brief Read a range if it is still cached.
+  ///
+  /// A cache miss returns an empty optional. I/O errors are propagated.
+  Result<std::optional<std::shared_ptr<Buffer>>> ReadIfCached(ReadRange range);
+
   /// \brief Wait until all ranges added so far have been cached.
   Future<> Wait();
 
   /// \brief Wait until all given ranges have been cached.
   Future<> WaitFor(std::vector<ReadRange> ranges);
 
+  /// \brief Evict cache entries ending at or before `end_offset`.
+  ///
+  /// An entry straddling `end_offset` is retained. Buffers already returned by
+  /// Read() stay valid through shared ownership.
+  void EvictEntriesBefore(int64_t end_offset);
+
  protected:
   struct Impl;
   struct LazyImpl;