wgtmac commented on code in PR #49855:
URL: https://github.com/apache/arrow/pull/49855#discussion_r4357464403
##########
cpp/src/arrow/io/caching.h:
##########
@@ -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);
Review Comment:
`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
##########
cpp/src/parquet/arrow/reader.cc:
##########
@@ -1325,10 +1336,27 @@
FileReaderImpl::GetRecordBatchGenerator(std::shared_ptr<FileReader> reader,
reader_properties_.cache_options());
END_PARQUET_CATCH_EXCEPTIONS
}
+ std::vector<int64_t> evict_before_offsets;
+ if (reader_properties_.pre_buffer() &&
reader_properties_.auto_evict_read_cache() &&
+ !column_indices.empty() && !row_group_indices.empty()) {
+ const int64_t kNoMoreRanges = std::numeric_limits<int64_t>::max();
+ evict_before_offsets.resize(row_group_indices.size() + 1, kNoMoreRanges);
+ 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)]},
Review Comment:
GetReadRanges has been called in PreBuffer already. But I don't have a good
answer to reuse result from PreBuffer since it needs a larger change which
might look ugly anyway. So I'm fine to keep it as is if it only affects the
eviction code path.
##########
cpp/src/parquet/file_reader.cc:
##########
@@ -617,8 +626,7 @@ class SerializedFile : public ParquetFileReader::Contents {
ReaderProperties properties_;
std::shared_ptr<PageIndexReader> page_index_reader_;
std::unique_ptr<BloomFilterReader> bloom_filter_reader_;
- // Maps row group ordinal and prebuffer status of its column chunks in the
form of a
- // bitmap buffer.
+
Review Comment:
This change seems unrelated?
##########
cpp/src/parquet/arrow/reader.cc:
##########
@@ -1207,12 +1208,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::vector<int64_t> evict_before_offsets)
: 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),
+ evict_before_offsets_(std::move(evict_before_offsets)),
Review Comment:
nit: add a DCHECK for evict_before_offsets_.size() just in case?
##########
cpp/src/parquet/arrow/arrow_reader_writer_test.cc:
##########
@@ -2726,6 +2728,201 @@ TEST(TestArrowReadWrite, GetRecordBatchReaderNoColumns)
{
ASSERT_EQ(actual_batch->num_rows(), num_rows);
}
+// 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);
+ const int num_rows = 1024;
+ const int row_group_size = 256;
+ const int num_columns = 3;
+ const std::vector<int> row_groups = {0, 1, 2, 3};
+ const std::vector<int> column_indices = {0, 1, 2};
+
+ std::shared_ptr<Table> table;
+ ASSERT_NO_FATAL_FAILURE(MakeDoubleTable(num_columns, num_rows, 1, &table));
+
+ std::shared_ptr<Buffer> buffer;
+ ASSERT_NO_FATAL_FAILURE(WriteTableToBuffer(table, row_group_size,
+
default_arrow_writer_properties(), &buffer));
+
+ std::unique_ptr<FileReader> reader;
+ FileReaderBuilder builder;
+ ASSERT_OK(builder.Open(std::make_shared<BufferReader>(buffer)));
+ ASSERT_OK(builder.properties(properties)->Build(&reader));
+ ASSERT_EQ(reader->num_row_groups(), static_cast<int>(row_groups.size()));
+
+ 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());
+
+ std::shared_ptr<Table> rg_table;
+ ASSERT_OK(reader->ReadRowGroup(/*i=*/0, column_indices, &rg_table));
+ ASSERT_EQ(rg_table->num_rows(), row_group_size);
+
+ // 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<int64_t>::max();
+ for (const auto& r : rg1_ranges) rg1_min = std::min(rg1_min, r.offset);
Review Comment:
Can we simply use `RowGroupMetadata::file_offset` instead?
##########
cpp/src/parquet/arrow/arrow_reader_writer_test.cc:
##########
@@ -2726,6 +2728,201 @@ TEST(TestArrowReadWrite, GetRecordBatchReaderNoColumns)
{
ASSERT_EQ(actual_batch->num_rows(), num_rows);
}
+// 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);
+ const int num_rows = 1024;
+ const int row_group_size = 256;
+ const int num_columns = 3;
+ const std::vector<int> row_groups = {0, 1, 2, 3};
+ const std::vector<int> column_indices = {0, 1, 2};
+
+ std::shared_ptr<Table> table;
+ ASSERT_NO_FATAL_FAILURE(MakeDoubleTable(num_columns, num_rows, 1, &table));
+
+ std::shared_ptr<Buffer> buffer;
+ ASSERT_NO_FATAL_FAILURE(WriteTableToBuffer(table, row_group_size,
+
default_arrow_writer_properties(), &buffer));
+
+ std::unique_ptr<FileReader> reader;
+ FileReaderBuilder builder;
+ ASSERT_OK(builder.Open(std::make_shared<BufferReader>(buffer)));
+ ASSERT_OK(builder.properties(properties)->Build(&reader));
+ ASSERT_EQ(reader->num_row_groups(), static_cast<int>(row_groups.size()));
+
+ 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());
+
+ std::shared_ptr<Table> rg_table;
+ ASSERT_OK(reader->ReadRowGroup(/*i=*/0, column_indices, &rg_table));
+ ASSERT_EQ(rg_table->num_rows(), row_group_size);
+
+ // 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<int64_t>::max();
+ for (const auto& r : rg1_ranges) rg1_min = std::min(rg1_min, r.offset);
+
+ reader->parquet_reader()->EvictPreBufferedDataBefore(rg1_min);
+ ASSERT_RAISES(Invalid,
+ reader->parquet_reader()->WhenBuffered({0},
column_indices).status());
+ ASSERT_OK(reader->ReadRowGroup(/*i=*/1, column_indices, &rg_table));
+ ASSERT_EQ(rg_table->num_rows(), row_group_size);
+
+ reader->parquet_reader()->EvictPreBufferedDataBefore(
+ std::numeric_limits<int64_t>::max());
+ ASSERT_RAISES(Invalid,
+ reader->parquet_reader()->WhenBuffered({1},
column_indices).status());
+ std::shared_ptr<ChunkedArray> 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.
+ reader->parquet_reader()->EvictPreBufferedDataBefore(rg1_min);
+
+ // A reader that never called PreBuffer is a no-op.
+ std::unique_ptr<FileReader> no_prebuffer_reader;
+ FileReaderBuilder no_prebuffer_builder;
+ ASSERT_OK(no_prebuffer_builder.Open(std::make_shared<BufferReader>(buffer)));
+ ASSERT_OK(no_prebuffer_builder.Build(&no_prebuffer_reader));
+ 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<int> row_groups = {0, 1, 2, 3};
+ const std::vector<int> column_indices = {0, 1};
+
+ std::shared_ptr<Table> table;
+ ASSERT_NO_FATAL_FAILURE(MakeDoubleTable(num_columns, num_rows, 1, &table));
+
+ std::shared_ptr<Buffer> buffer;
+ ASSERT_NO_FATAL_FAILURE(WriteTableToBuffer(table, row_group_size,
+
default_arrow_writer_properties(), &buffer));
+
+ std::unique_ptr<FileReader> reader;
+ FileReaderBuilder builder;
+ ASSERT_OK(builder.Open(std::make_shared<BufferReader>(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<int64_t>(buffer->size());
+ options.range_size_limit = static_cast<int64_t>(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());
+
+ // 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<int64_t>::max();
+ for (const auto& r : rg3_ranges) rg3_min = std::min(rg3_min, r.offset);
+ reader->parquet_reader()->EvictPreBufferedDataBefore(rg3_min);
+ ASSERT_OK(reader->parquet_reader()->WhenBuffered(row_groups,
column_indices).status());
+
+ // An offset past the whole file frees the spanning entry.
+ reader->parquet_reader()->EvictPreBufferedDataBefore(
+ static_cast<int64_t>(buffer->size()));
+ ASSERT_RAISES(
+ Invalid,
+ reader->parquet_reader()->WhenBuffered(row_groups,
column_indices).status());
+}
+
+// GH-39808: when Dataset.to_batches-style iteration drives the async
Review Comment:
I'm not sure if it is better to use a real dataset reader test case to
replace this one.
##########
cpp/src/parquet/arrow/arrow_reader_writer_test.cc:
##########
@@ -2726,6 +2728,201 @@ TEST(TestArrowReadWrite, GetRecordBatchReaderNoColumns)
{
ASSERT_EQ(actual_batch->num_rows(), num_rows);
}
+// 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);
+ const int num_rows = 1024;
+ const int row_group_size = 256;
+ const int num_columns = 3;
+ const std::vector<int> row_groups = {0, 1, 2, 3};
+ const std::vector<int> column_indices = {0, 1, 2};
+
+ std::shared_ptr<Table> table;
+ ASSERT_NO_FATAL_FAILURE(MakeDoubleTable(num_columns, num_rows, 1, &table));
+
+ std::shared_ptr<Buffer> buffer;
+ ASSERT_NO_FATAL_FAILURE(WriteTableToBuffer(table, row_group_size,
+
default_arrow_writer_properties(), &buffer));
+
+ std::unique_ptr<FileReader> reader;
+ FileReaderBuilder builder;
+ ASSERT_OK(builder.Open(std::make_shared<BufferReader>(buffer)));
+ ASSERT_OK(builder.properties(properties)->Build(&reader));
+ ASSERT_EQ(reader->num_row_groups(), static_cast<int>(row_groups.size()));
+
+ 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());
+
+ std::shared_ptr<Table> rg_table;
+ ASSERT_OK(reader->ReadRowGroup(/*i=*/0, column_indices, &rg_table));
Review Comment:
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.
##########
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);
Review Comment:
I think we need at least one test case for the dataset reader.
##########
cpp/src/arrow/io/caching.cc:
##########
@@ -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)
Review Comment:
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.
##########
cpp/src/arrow/io/caching.cc:
##########
@@ -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.
Review Comment:
```suggestion
// Prefetch immediately, regardless of executor availability, if possible
```
I would suggest reverting this comment change since it is too verbose.
##########
cpp/src/arrow/io/memory_test.cc:
##########
@@ -918,6 +940,155 @@ TEST(RangeReadCache, LazyWithPrefetching) {
ASSERT_RAISES(Invalid, cache.Read({25, 2}));
}
+TEST(RangeReadCache, EvictEntriesBefore) {
+ // GH-39808: entries cached by PreBuffer()-style code need to be evictable so
+ // memory stays bounded while iterating a large Parquet file.
+ std::string data = "abcdefghijklmnopqrstuvwxyz";
+
+ for (auto lazy : std::vector<bool>{false, true}) {
+ SCOPED_TRACE(lazy);
+ CacheOptions options = CacheOptions::Defaults();
+ options.hole_size_limit = 0; // disable coalescing: one entry per range
+ options.range_size_limit = 10;
+ options.lazy = lazy;
+
+ auto file = std::make_shared<BufferReader>(std::make_shared<Buffer>(data));
+ internal::ReadRangeCache cache(file, {}, options);
+
+ // Entries: [1,3), [10,14), [20,22).
+ ASSERT_OK(cache.Cache({{1, 2}, {10, 4}, {20, 2}}));
+ ASSERT_OK_AND_ASSIGN(auto buf, cache.Read({10, 4}));
+ AssertBufferEqual(*buf, "klmn");
+
+ // An offset that splits no entry frees nothing.
+ cache.EvictEntriesBefore(0);
+ cache.EvictEntriesBefore(2); // [1,3) extends past 2 -> retained
+ ASSERT_OK(cache.Read({1, 2}));
+
+ // end_offset == an entry's end frees exactly that entry ([1,3) ends at 3).
+ 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.
+ cache.EvictEntriesBefore(12); // [10,14) straddles 12
+ ASSERT_OK(cache.Read({10, 4}));
+
+ // A wide offset frees every remaining entry in one call.
+ cache.EvictEntriesBefore(100);
+ ASSERT_RAISES(Invalid, cache.Read({10, 4}));
+ ASSERT_RAISES(Invalid, cache.Read({20, 2}));
+
+ // Empty cache is a safe no-op.
+ cache.EvictEntriesBefore(100);
+ }
+}
+
+TEST(RangeReadCache, ConcurrentReadAndEvict) {
+ // 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
+ // 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<bool>{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<BufferReader>(std::make_shared<Buffer>(data));
+ internal::ReadRangeCache cache(file, {}, options);
+
+ std::vector<ReadRange> ranges;
+ ranges.reserve(kNumRanges);
+ for (int i = 0; i < kNumRanges; ++i) {
+ ranges.push_back({static_cast<int64_t>(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<bool> stop{false};
+ std::atomic<int> 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 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<ReadRange> 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<std::thread> threads;
Review Comment:
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.
##########
cpp/src/parquet/arrow/arrow_reader_writer_test.cc:
##########
@@ -2726,6 +2728,201 @@ TEST(TestArrowReadWrite, GetRecordBatchReaderNoColumns)
{
ASSERT_EQ(actual_batch->num_rows(), num_rows);
}
+// 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);
+ const int num_rows = 1024;
+ const int row_group_size = 256;
+ const int num_columns = 3;
+ const std::vector<int> row_groups = {0, 1, 2, 3};
+ const std::vector<int> column_indices = {0, 1, 2};
+
+ std::shared_ptr<Table> table;
+ ASSERT_NO_FATAL_FAILURE(MakeDoubleTable(num_columns, num_rows, 1, &table));
+
+ std::shared_ptr<Buffer> buffer;
+ ASSERT_NO_FATAL_FAILURE(WriteTableToBuffer(table, row_group_size,
+
default_arrow_writer_properties(), &buffer));
+
+ std::unique_ptr<FileReader> reader;
+ FileReaderBuilder builder;
+ ASSERT_OK(builder.Open(std::make_shared<BufferReader>(buffer)));
+ ASSERT_OK(builder.properties(properties)->Build(&reader));
+ ASSERT_EQ(reader->num_row_groups(), static_cast<int>(row_groups.size()));
+
+ 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());
+
+ std::shared_ptr<Table> rg_table;
+ ASSERT_OK(reader->ReadRowGroup(/*i=*/0, column_indices, &rg_table));
+ ASSERT_EQ(rg_table->num_rows(), row_group_size);
+
+ // 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<int64_t>::max();
+ for (const auto& r : rg1_ranges) rg1_min = std::min(rg1_min, r.offset);
+
+ reader->parquet_reader()->EvictPreBufferedDataBefore(rg1_min);
+ ASSERT_RAISES(Invalid,
+ reader->parquet_reader()->WhenBuffered({0},
column_indices).status());
+ ASSERT_OK(reader->ReadRowGroup(/*i=*/1, column_indices, &rg_table));
+ ASSERT_EQ(rg_table->num_rows(), row_group_size);
+
+ reader->parquet_reader()->EvictPreBufferedDataBefore(
+ std::numeric_limits<int64_t>::max());
+ ASSERT_RAISES(Invalid,
+ reader->parquet_reader()->WhenBuffered({1},
column_indices).status());
+ std::shared_ptr<ChunkedArray> 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.
+ reader->parquet_reader()->EvictPreBufferedDataBefore(rg1_min);
+
+ // A reader that never called PreBuffer is a no-op.
+ std::unique_ptr<FileReader> no_prebuffer_reader;
+ FileReaderBuilder no_prebuffer_builder;
+ ASSERT_OK(no_prebuffer_builder.Open(std::make_shared<BufferReader>(buffer)));
+ ASSERT_OK(no_prebuffer_builder.Build(&no_prebuffer_reader));
+ 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<int> row_groups = {0, 1, 2, 3};
+ const std::vector<int> column_indices = {0, 1};
+
+ std::shared_ptr<Table> table;
+ ASSERT_NO_FATAL_FAILURE(MakeDoubleTable(num_columns, num_rows, 1, &table));
+
+ std::shared_ptr<Buffer> buffer;
+ ASSERT_NO_FATAL_FAILURE(WriteTableToBuffer(table, row_group_size,
+
default_arrow_writer_properties(), &buffer));
+
+ std::unique_ptr<FileReader> reader;
+ FileReaderBuilder builder;
+ ASSERT_OK(builder.Open(std::make_shared<BufferReader>(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<int64_t>(buffer->size());
+ options.range_size_limit = static_cast<int64_t>(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());
+
+ // 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<int64_t>::max();
+ for (const auto& r : rg3_ranges) rg3_min = std::min(rg3_min, r.offset);
+ reader->parquet_reader()->EvictPreBufferedDataBefore(rg3_min);
+ ASSERT_OK(reader->parquet_reader()->WhenBuffered(row_groups,
column_indices).status());
+
+ // An offset past the whole file frees the spanning entry.
+ reader->parquet_reader()->EvictPreBufferedDataBefore(
+ static_cast<int64_t>(buffer->size()));
+ ASSERT_RAISES(
+ Invalid,
+ reader->parquet_reader()->WhenBuffered(row_groups,
column_indices).status());
+}
+
+// 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> table;
+ ASSERT_NO_FATAL_FAILURE(MakeDoubleTable(num_columns, num_rows, 1, &table));
+
+ std::shared_ptr<Buffer> buffer;
+ ASSERT_NO_FATAL_FAILURE(WriteTableToBuffer(table, row_group_size,
+
default_arrow_writer_properties(), &buffer));
+
+ // General FileReaders retain their cache because callers may create another
+ // generator for the same row group.
+ ASSERT_FALSE(properties.auto_evict_read_cache());
+ std::shared_ptr<FileReader> retaining_reader;
+ {
+ std::unique_ptr<FileReader> unique_reader;
+ FileReaderBuilder builder;
+ ASSERT_OK(builder.Open(std::make_shared<BufferReader>(buffer)));
+ ASSERT_OK(builder.properties(properties)->Build(&unique_reader));
+ retaining_reader = std::move(unique_reader);
+ }
+ ASSERT_OK_AND_ASSIGN(
+ auto retaining_generator,
+ retaining_reader->GetRecordBatchGenerator(retaining_reader, {0}, {0}));
+ auto retained_fut = retaining_generator();
+ ASSERT_OK_AND_ASSIGN(auto retained_batch, retained_fut.result());
+ ASSERT_NE(retained_batch, nullptr);
+ auto retained_end_fut = retaining_generator();
+ ASSERT_OK_AND_ASSIGN(auto retained_end, retained_end_fut.result());
+ ASSERT_EQ(retained_end, nullptr);
+ ASSERT_OK(retaining_reader->parquet_reader()->WhenBuffered({0},
{0}).status());
+
+ properties.set_auto_evict_read_cache(true);
+
+ std::shared_ptr<FileReader> reader;
+ {
+ std::unique_ptr<FileReader> unique_reader;
+ FileReaderBuilder builder;
+ ASSERT_OK(builder.Open(std::make_shared<BufferReader>(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},
+ /*cpu_executor=*/nullptr,
+ /*rows_to_readahead=*/2 *
row_group_size));
+ std::vector<std::shared_ptr<::arrow::RecordBatch>> 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);
+ ASSERT_RAISES(Invalid,
Review Comment:
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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
