justinli500 commented on code in PR #49855:
URL: https://github.com/apache/arrow/pull/49855#discussion_r4520955573
##########
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 added the Dataset test and kept the FileReader test as well. The
FileReader test still covers default cache retention and explicit eviction.
--
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]
