wgtmac commented on code in PR #49855:
URL: https://github.com/apache/arrow/pull/49855#discussion_r4547906867
##########
cpp/src/arrow/io/memory_test.cc:
##########
@@ -810,21 +842,24 @@ TEST(RangeReadCache, Concurrency) {
{
internal::ReadRangeCache cache(file, {}, options);
ASSERT_OK(cache.Cache(ranges));
- std::vector<Future<std::shared_ptr<Buffer>>> futures;
+ std::vector<Future<std::optional<std::shared_ptr<Buffer>>>> futures;
for (const auto& range : ranges) {
futures.push_back(
cache.WaitFor({range}).Then([&cache, range]() { return
cache.Read(range); }));
}
for (auto fut : futures) {
- ASSERT_FINISHES_OK(fut);
+ ASSERT_OK_AND_ASSIGN(auto buffer, fut.result());
Review Comment:
Please preserve the bounded wait when unpacking the future, for example with
`ASSERT_FINISHES_OK_AND_ASSIGN(auto buffer, fut)`. Calling `fut.result()`
directly removes the existing timeout.
##########
cpp/src/arrow/io/memory_test.cc:
##########
@@ -761,29 +774,37 @@ TEST(RangeReadCache, Basics) {
ASSERT_OK(cache.Cache({{10, 4}, {14, 0}, {15, 4}}));
ASSERT_OK_AND_ASSIGN(auto buf, cache.Read({20, 2}));
- AssertBufferEqual(*buf, "uv");
+ ASSERT_TRUE(buf);
+ AssertBufferEqual(**buf, "uv");
Review Comment:
There seems to be a lot of duplication. Perhaps we can use a lambda to save
some lines.
##########
cpp/src/arrow/io/memory_test.cc:
##########
@@ -886,36 +925,247 @@ TEST(RangeReadCache, LazyWithPrefetching) {
ASSERT_EQ(0, file->read_count());
ASSERT_OK_AND_ASSIGN(auto buf, cache.Read({8, 2}));
- AssertBufferEqual(*buf, "ij");
+ ASSERT_TRUE(buf);
+ AssertBufferEqual(**buf, "ij");
// Read {8, 2} and prefetch {20, 2}
ASSERT_EQ(2, file->read_count());
ASSERT_OK_AND_ASSIGN(buf, cache.Read({20, 2}));
- AssertBufferEqual(*buf, "uv");
+ ASSERT_TRUE(buf);
+ AssertBufferEqual(**buf, "uv");
// Read count remains 2 as the range {20, 2} has already been prefetched
ASSERT_EQ(2, file->read_count());
ASSERT_OK_AND_ASSIGN(buf, cache.Read({1, 1}));
- AssertBufferEqual(*buf, "b");
+ ASSERT_TRUE(buf);
+ AssertBufferEqual(**buf, "b");
// Read {1, 3} and prefetch {5, 2}
ASSERT_EQ(4, file->read_count());
ASSERT_OK_AND_ASSIGN(buf, cache.Read({3, 1}));
- AssertBufferEqual(*buf, "d");
+ ASSERT_TRUE(buf);
+ AssertBufferEqual(**buf, "d");
// Already prefetched
ASSERT_EQ(4, file->read_count());
// Requested ranges are still cached
ASSERT_OK_AND_ASSIGN(buf, cache.Read({5, 1}));
- AssertBufferEqual(*buf, "f");
+ ASSERT_TRUE(buf);
+ AssertBufferEqual(**buf, "f");
// Already prefetched
ASSERT_EQ(4, file->read_count());
// Non-cached ranges
- ASSERT_RAISES(Invalid, cache.Read({20, 3}));
- ASSERT_RAISES(Invalid, cache.Read({19, 3}));
- ASSERT_RAISES(Invalid, cache.Read({0, 3}));
- ASSERT_RAISES(Invalid, cache.Read({25, 2}));
+ ASSERT_OK_AND_EQ(std::nullopt, cache.Read({20, 3}));
+ ASSERT_OK_AND_EQ(std::nullopt, cache.Read({19, 3}));
+ ASSERT_OK_AND_EQ(std::nullopt, cache.Read({0, 3}));
+ ASSERT_OK_AND_EQ(std::nullopt, 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}));
+ ASSERT_TRUE(buf);
+ 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_AND_ASSIGN(buf, cache.Read({1, 2}));
+ ASSERT_TRUE(buf);
+
+ // end_offset == an entry's end frees exactly that entry ([1,3) ends at 3).
+ cache.EvictEntriesBefore(3);
+ ASSERT_OK_AND_EQ(std::nullopt, cache.Read({1, 2}));
+ ASSERT_OK_AND_ASSIGN(buf, cache.Read({10, 4})); // others intact
+ ASSERT_TRUE(buf);
+ AssertBufferEqual(**buf, "klmn");
+
+ // An offset inside an entry leaves it in place.
+ cache.EvictEntriesBefore(12); // [10,14) straddles 12
+ ASSERT_OK_AND_ASSIGN(buf, cache.Read({10, 4}));
+ ASSERT_TRUE(buf);
+
+ // A wide offset frees every remaining entry in one call.
+ cache.EvictEntriesBefore(100);
+ ASSERT_OK_AND_EQ(std::nullopt, cache.Read({10, 4}));
+ ASSERT_OK_AND_EQ(std::nullopt, cache.Read({20, 2}));
+ AssertBufferEqual(**buf, "klmn");
+
+ // Empty cache is a safe no-op.
+ cache.EvictEntriesBefore(100);
+ }
+}
+
+TEST(RangeReadCache, EvictAndRecache) {
+ for (bool lazy : {false, true}) {
+ SCOPED_TRACE(lazy);
+ auto options = CacheOptions::Defaults();
+ options.hole_size_limit = 0;
+ options.range_size_limit = 1;
+ options.lazy = lazy;
+ options.prefetch_limit = 1;
+ auto file = std::make_shared<CountingBufferReader>(
+ Buffer::FromString("abcdefghijklmnopqrstuvwxyz"));
+ internal::ReadRangeCache cache(file, {}, options);
+ ASSERT_OK(cache.Cache({{1, 2}, {10, 2}, {20, 2}}));
+
+ cache.EvictEntriesBefore(3);
+ ASSERT_FINISHES_OK(cache.Wait());
+ const int64_t expected_reads = lazy ? 2 : 3;
+ ASSERT_EQ(file->read_count(), expected_reads);
+ ASSERT_OK_AND_EQ(std::nullopt, cache.Read({1, 2}));
+ ASSERT_FINISHES_AND_RAISES(Invalid, cache.WaitFor({{1, 2}}));
+
+ ASSERT_OK(cache.Cache({{1, 2}}));
+ ASSERT_FINISHES_OK(cache.WaitFor({{1, 2}}));
+ ASSERT_OK_AND_ASSIGN(auto buffer, cache.Read({1, 2}));
+ ASSERT_TRUE(buffer);
+ AssertBufferEqual(**buffer, "bc");
+ ASSERT_EQ(file->read_count(), expected_reads + 1);
+
+ cache.EvictEntriesBefore(12);
+ ASSERT_FINISHES_OK(cache.Wait());
+ ASSERT_OK_AND_EQ(std::nullopt, cache.Read({1, 2}));
+ ASSERT_OK_AND_EQ(std::nullopt, cache.Read({10, 2}));
+ ASSERT_OK_AND_ASSIGN(buffer, cache.Read({20, 2}));
+ ASSERT_TRUE(buffer);
+ AssertBufferEqual(**buffer, "uv");
+ ASSERT_EQ(file->read_count(), expected_reads + 1);
+
+ cache.EvictEntriesBefore(100);
+ ASSERT_FINISHES_OK(cache.Wait());
+ ASSERT_OK(cache.Cache({{5, 2}}));
+ ASSERT_FINISHES_OK(cache.Wait());
+ ASSERT_OK_AND_ASSIGN(buffer, cache.Read({5, 2}));
+ ASSERT_TRUE(buffer);
+ AssertBufferEqual(**buffer, "fg");
+ ASSERT_EQ(file->read_count(), expected_reads + 2);
+ }
+}
+
+TEST(RangeReadCache, ConcurrentReadAndEvict) {
+#ifndef ARROW_ENABLE_THREADING
+ GTEST_SKIP() << "Test requires threading support";
+#endif
+ // 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 = [&]() {
Review Comment:
Could we coordinate the readers and evictor and assert that reads actually
overlap the eviction loop? The evictor can set `stop` before any reader runs.
Distinct data per range would also catch wrong-range reads.
##########
cpp/src/parquet/arrow/arrow_reader_writer_test.cc:
##########
@@ -2726,6 +2728,214 @@ 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();
+ // Control pre-buffering explicitly so row-group reads preserve the cache.
+ properties.set_pre_buffer(false);
+ 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()));
+
+ auto options = ::arrow::io::CacheOptions::LazyDefaults();
+ options.hole_size_limit = 0;
+ options.range_size_limit = 1;
+ reader->parquet_reader()->PreBuffer(row_groups, column_indices,
+ ::arrow::io::IOContext(), options);
+ 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));
+ AssertTablesEqual(*table->Slice(0, row_group_size), *rg_table,
+ /*same_chunk_layout=*/false);
+ ASSERT_OK(reader->parquet_reader()->WhenBuffered(row_groups,
column_indices).status());
+
+ const int64_t rg1_min =
+ reader->parquet_reader()->metadata()->RowGroup(1)->file_offset();
+
+ reader->parquet_reader()->EvictPreBufferedDataBefore(rg1_min);
+ ASSERT_RAISES(Invalid,
+ reader->parquet_reader()->WhenBuffered({0},
column_indices).status());
+ ASSERT_OK(reader->parquet_reader()->WhenBuffered({1, 2, 3},
column_indices).status());
+ ASSERT_OK(reader->ReadRowGroup(/*i=*/1, column_indices, &rg_table));
+ AssertTablesEqual(*table->Slice(row_group_size, row_group_size), *rg_table,
+ /*same_chunk_layout=*/false);
+
+ 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));
+ ::arrow::AssertChunkedEquivalent(
+ *table->column(0)->Slice(row_group_size, row_group_size),
*reread_column);
+
+ // 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());
+}
+
+TEST(TestArrowReadWrite, GetRecordBatchGeneratorReleasesPreBufferedRowGroups) {
+ ArrowReaderProperties properties = default_arrow_reader_properties();
+ properties.set_pre_buffer(true);
+ 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);
+ properties.set_batch_size(row_group_size / 2);
+ auto options = ::arrow::io::CacheOptions::Defaults();
+ options.hole_size_limit = 0;
+ options.range_size_limit = 1;
+ properties.set_cache_options(options);
+
+ 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);
+
+ ASSERT_OK_AND_ASSIGN(
+ auto batch_generator,
+ reader->GetRecordBatchGenerator(reader, {0, 1, 2, 3}, {0, 1},
Review Comment:
Could we add an out-of-order case such as `{2, 0, 3, 1}`, with no readahead,
and check that upcoming groups remain cached? The ascending case also passes if
the suffix minimum is incorrectly replaced by the next row group's offset.
--
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]
