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/dataset/file_parquet_test.cc

@@ -32,8 +32,10 @@
 #include "arrow/io/memory.h"
 #include "arrow/io/test_common.h"
 #include "arrow/io/util_internal.h"
+#include "arrow/memory_pool.h"
 #include "arrow/record_batch.h"
 #include "arrow/table.h"
+#include "arrow/testing/executor_util.h"
 #include "arrow/testing/gtest_util.h"
 #include "arrow/testing/util.h"
 #include "arrow/type.h"
@@ -375,6 +377,82 @@ TEST_F(TestParquetFileFormat, CountRowsPredicatePushdown) {
   }
 }
 
+TEST_F(TestParquetFileFormat, ScanBatchesAsyncReleasesPreBufferedRowGroups) {
+  class CopyingBufferReader : public io::BufferReader {
+   public:
+    using io::BufferReader::BufferReader;
+
+    Future<std::shared_ptr<Buffer>> ReadAsync(const io::IOContext& ctx, int64_t position,
+                                              int64_t nbytes,
+                                              bool allow_short_read) override {
+      return io::BufferReader::ReadAsync(ctx, position, nbytes, allow_short_read)
+          .Then([pool = ctx.pool()](const std::shared_ptr<Buffer>& buffer) {
+            return buffer->CopySlice(0, buffer->size(), pool);
+          });
+    }
+  };
+
+  SetSchema({field("i32", int32())});
+  SetFilter(literal(true));
+  auto expected = RecordBatchFromJSON(opts_->dataset_schema,
+                                      "[[0], [1], [2], [3], [4], [5], [6], [7], "
+                                      "[8], [9], [10], [11], [12], [13], [14], [15]]");
+  constexpr int kNumRowGroups = 4;
+  constexpr int kRowGroupSize = 4;
+  RecordBatchVector row_groups;
+  for (int i = 0; i < kNumRowGroups; ++i) {
+    row_groups.push_back(expected->Slice(i * kRowGroupSize, kRowGroupSize));
+  }
+  ASSERT_OK_AND_ASSIGN(auto reader, RecordBatchReader::Make(std::move(row_groups)));
+  ASSERT_OK_AND_ASSIGN(auto buffer, ParquetFormatHelper::Write(reader.get()));
+
+  // Copy source reads into a separate pool to track only cached compressed bytes.
+  ProxyMemoryPool read_pool(default_memory_pool());
+  MockExecutor executor;
+  auto input = std::make_shared<CopyingBufferReader>(buffer);
+  auto fragment = MakeFragment(FileSource(input, buffer->size()));
+  ASSERT_OK(
+      checked_pointer_cast<ParquetFileFragment>(fragment)->EnsureCompleteMetadata());
+
+  auto scan_options = std::make_shared<ParquetFragmentScanOptions>();
+  scan_options->arrow_reader_properties->set_pre_buffer(true);
+  scan_options->arrow_reader_properties->set_io_context(io::IOContext(&read_pool));
+  auto cache_options = io::CacheOptions::Defaults();
+  cache_options.hole_size_limit = 0;
+  cache_options.range_size_limit = 1;
+  scan_options->arrow_reader_properties->set_cache_options(cache_options);
+  opts_->fragment_scan_options = scan_options;
+  opts_->batch_size = kRowGroupSize / 2;
+  opts_->batch_readahead = 0;
+  opts_->use_threads = false;
+  opts_->cpu_executor = &executor;
+
+  ASSERT_OK_AND_ASSIGN(auto generator, fragment->ScanBatchesAsync(opts_));
+  int64_t cached_bytes = read_pool.bytes_allocated();
+  ASSERT_GT(cached_bytes, 0);
+  for (int i = 0; i < kNumRowGroups * 2; ++i) {
+    auto next = generator();
+    ASSERT_OK_AND_ASSIGN(auto batch, next.result());
+    ASSERT_NE(batch, nullptr);
+    AssertBatchesEqual(*expected->Slice(i * opts_->batch_size, opts_->batch_size),
+                       *batch);
+    // The concatenating generator's queued batch stays within this row group.
+    if (i % 2 == 0) {
+      const int64_t remaining_bytes = read_pool.bytes_allocated();
+      ASSERT_LT(remaining_bytes, cached_bytes);
+      if (i / 2 + 1 < kNumRowGroups) {
+        ASSERT_GT(remaining_bytes, 0);
+      } else {
+        ASSERT_EQ(remaining_bytes, 0);
+      }
+      cached_bytes = remaining_bytes;
+    }
+  }
+  auto end = generator();
+  ASSERT_OK_AND_ASSIGN(auto batch, end.result());
+  ASSERT_EQ(batch, nullptr);
+}
+
 TEST_F(TestParquetFileFormat, CachedMetadata) {
   // Create a test file
   auto mock_fs = std::make_shared<fs::internal::MockFileSystem>(fs::kNoTime);

cpp/src/arrow/filesystem/s3fs_benchmark.cc

@@ -258,7 +258,6 @@ static void CoalescedRead(benchmark::State& st, S3FileSystem* fs,
   int total_items = 0;
   for (auto _ : st) {
     std::shared_ptr<io::RandomAccessFile> file;
-    std::shared_ptr<Buffer> buf;
     int64_t size = 0;
     ASSERT_OK_AND_ASSIGN(file, fs->OpenInputFile(path));
     ASSERT_OK_AND_ASSIGN(size, file->GetSize());
@@ -281,8 +280,9 @@ static void CoalescedRead(benchmark::State& st, S3FileSystem* fs,
     offset = 0;
     while (offset < size) {
       const int64_t read = std::min(size, kChunkSize);
-      ASSERT_OK_AND_ASSIGN(buf, cache.Read({offset, read}));
-      total_bytes += buf->size();
+      ASSERT_OK_AND_ASSIGN(auto buf, cache.Read({offset, read}));
+      ASSERT_TRUE(buf);
+      total_bytes += (*buf)->size();
       offset += read;
     }
   }

cpp/src/arrow/io/caching.cc

@@ -18,7 +18,9 @@
 #include <algorithm>
 #include <atomic>
 #include <cmath>
+#include <iterator>
 #include <mutex>
+#include <optional>
 #include <utility>
 #include <vector>
 
@@ -151,12 +153,13 @@ 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::vector<RangeCacheEntry> entries;  // GUARDED_BY(entry_mutex)
+  size_t first_live = 0;                 // 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,59 @@ 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);
+      if (first_live == entries.size()) {
+        entries = std::move(new_entries);
+      } else {
+        std::vector<RangeCacheEntry> merged(entries.size() - first_live +
+                                            new_entries.size());
+        std::merge(std::make_move_iterator(entries.begin() + first_live),
+                   std::make_move_iterator(entries.end()),
+                   std::make_move_iterator(new_entries.begin()),
+                   std::make_move_iterator(new_entries.end()), merged.begin());
+        entries = std::move(merged);
+      }
+      first_live = 0;
     }
     // Prefetch immediately, regardless of executor availability, if possible
-    auto st = file->WillNeed(ranges);
+    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>>> Read(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() + first_live, 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 +242,73 @@ 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));
   }
 
-  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() - first_live);
+      for (auto it = entries.begin() + first_live; it != entries.end(); ++it) {
+        futures.emplace_back(MaybeRead(&*it));
+      }
     }
     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);
+    while (first_live < entries.size()) {
+      auto& entry = entries[first_live];
+      if (entry.range.offset + entry.range.length > end_offset) {
+        break;
+      }
+      entry.future = {};
+      ++first_live;
+    }
+  }
+
   // 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() + first_live, 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 +327,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,
@@ -329,7 +345,7 @@ Status ReadRangeCache::Cache(std::vector<ReadRange> ranges) {
   return impl_->Cache(std::move(ranges));
 }
 
-Result<std::shared_ptr<Buffer>> ReadRangeCache::Read(ReadRange range) {
+Result<std::optional<std::shared_ptr<Buffer>>> ReadRangeCache::Read(ReadRange range) {
   return impl_->Read(range);
 }
 
@@ -339,6 +355,10 @@ 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