You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by gi...@apache.org on 2026/10/08 17:42:22 UTC
(arrow-rs) branch gh-readonly-queue/main/pr-11235-9c07b17255ee65a70c1885368c1d91ca737521de created (now 890df9d6da)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-11235-9c07b17255ee65a70c1885368c1d91ca737521de
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
at 890df9d6da fix(parquet): release the bytes of a row group when the push decoder is done with it (#11235)
This branch includes the following new commits:
new 890df9d6da fix(parquet): release the bytes of a row group when the push decoder is done with it (#11235)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
(arrow-rs) 01/01: fix(parquet): release the bytes of a row group when the push decoder is done with it (#11235)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-11235-9c07b17255ee65a70c1885368c1d91ca737521de
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
commit 890df9d6da2df2c0cfba3b6547ea1fb89290c411
Author: Adrian Garcia Badaracco <17...@users.noreply.github.com>
AuthorDate: Thu Oct 8 17:41:57 2026 +0000
fix(parquet): release the bytes of a row group when the push decoder is done with it (#11235)
# Which issue does this PR close?
- Closes https://github.com/apache/arrow-rs/issues/11233.
- Part of https://github.com/apache/arrow-rs/issues/11234.
- closes https://github.com/apache/arrow-rs/issues/9695
- Replaces #9697
# Rationale for this change
`ParquetPushDecoder` does not release a pushed buffer that is larger
than a requested range. A caller that merges nearby ranges into one
object storage request thus holds all the bytes of the scan until the
scan ends.
| After row group | `buffered_bytes()` before this PR | With this PR |
|---|---:|---:|
| 0 | 195,128 | 0 |
| 1 | 390,256 | 0 |
| 2 | 585,384 | 0 |
| 3 | 780,512 | 0 |
The reproduction is in https://github.com/apache/arrow-rs/issues/11233.
# What changes are included in this PR?
1. When the decoder is done with a row group, it releases the pushed
bytes of all column chunks of that row group.
2. `into_builder().build()` releases the bytes that the rebuilt decoder
does not read, for example the bytes of row groups that a dynamic filter
pruned.
The inline comments show where to start and which parts of the diff are
mechanical.
<details>
<summary>Details: bytes that stay buffered, internals, performance
commits</summary>
These bytes stay buffered, as before, until `clear_all_ranges`:
- bytes that are not in a column chunk (for example, the footer in a
push of the full file), unless the decoder is rebuilt;
- bytes of row groups that the decoder skips during the scan (no
selected rows, or after the limit);
- bytes of a row group that the scan reads again later.
`buffered_bytes()` counts the bytes that the decoder holds. The kept
parts are zero-copy slices, so the allocator frees a pushed `Bytes` only
after all of its slices are dropped.
Internals: `PushBuffers::release_ranges` trims or splits pushed buffers.
`retain_ranges` calls `release_ranges` with the complement.
https://github.com/apache/arrow-rs/issues/11234 uses them again.
The commits after the fix make the bookkeeping cheaper (from the review
by @HippoBaro):
- the decoder counts queued row groups, and does not scan the queue for
each finished row group;
- `buffered_bytes()` is O(1);
- `release_ranges` trims buffers in place.
This PR includes https://github.com/apache/arrow-rs/pull/11236, which
changed the same code.
</details>
# Are these changes tested?
Yes. `test_decoder_releases_coalesced_push` is the reproduction from the
issue, and it fails without the change.
<details>
<summary>All new and changed tests</summary>
| Test | Case |
|---|---|
| `test_decoder_releases_coalesced_push` | the reproduction from the
issue |
| `test_decoder_releases_unrequested_bytes_of_row_group` | the pushed
buffer also holds a column that the scan does not read |
| `test_decoder_releases_coalesced_push_with_predicate` | a `RowFilter`,
which makes two requests per row group |
| `test_decoder_clear_all_ranges` (changed) | a push of the full file:
the bytes of each finished row group are released |
| `test_decoder_keeps_bytes_of_a_row_group_that_is_read_again` | a row
group that is read two times keeps its bytes |
| `test_into_builder_releases_bytes_of_skipped_row_groups` | read ahead
the full file, rebuild to read only row group 0 |
| `test_into_builder_preserves_buffered_bytes` (changed) | the rebuild
keeps the bytes of the next row group and releases the rest of the file
|
| `test_into_builder_keeps_bytes_of_predicate_columns` | a rebuild keeps
the bytes of a predicate column that is not in the output |
There are also unit tests for `release_ranges` (one compares it with a
rebuild of all buffers) and `retain_ranges`.
</details>
# Are there any user-facing changes?
There is no API change. There are two behavior changes:
- `buffered_bytes()` decreases after each row group when the caller
pushes buffers that are larger than the requested ranges.
- After `into_builder().build()` (or `build()` with `with_buffers`), the
decoder holds only the bytes that it reads. A caller that narrows the
projection and then widens it again must push the other columns again.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
---------
Co-authored-by: Claude Opus 5.5 <no...@anthropic.com>
---
parquet/src/arrow/push_decoder/mod.rs | 237 ++++++++++-
.../src/arrow/push_decoder/reader_builder/mod.rs | 74 +++-
parquet/src/arrow/push_decoder/remaining.rs | 28 +-
.../src/arrow/push_decoder/scan_plan/frontier.rs | 182 ++++++++-
parquet/src/util/push_buffers.rs | 433 ++++++++++++++++++++-
5 files changed, 933 insertions(+), 21 deletions(-)
diff --git a/parquet/src/arrow/push_decoder/mod.rs b/parquet/src/arrow/push_decoder/mod.rs
index 1f5054cc45..d2ef74e489 100644
--- a/parquet/src/arrow/push_decoder/mod.rs
+++ b/parquet/src/arrow/push_decoder/mod.rs
@@ -245,6 +245,10 @@ impl ParquetPushDecoderBuilder {
/// Provide a preexisting [`PushBuffers`] for the built decoder to read
/// from, so bytes already fetched are not requested again.
+ ///
+ /// [`Self::build`] releases the bytes that the decoder does not read:
+ /// the bytes outside the read column chunks of the row groups that the
+ /// decoder reads.
pub fn with_buffers(self, buffers: PushBuffers) -> Self {
Self {
input: PushDecoderInput { buffers },
@@ -347,7 +351,7 @@ impl ParquetPushDecoderBuilder {
);
// Initialize the decoder with the configured options
- let remaining_row_groups = RemainingRowGroups::new(
+ let mut remaining_row_groups = RemainingRowGroups::new(
schema,
parquet_metadata,
row_group_plan,
@@ -355,6 +359,10 @@ impl ParquetPushDecoderBuilder {
has_predicates,
row_group_reader_builder,
)?;
+ // The buffers can hold bytes that this decoder does not read. For
+ // example, the buffers of a rebuilt decoder can hold the bytes of row
+ // groups that the new configuration skips.
+ remaining_row_groups.release_unread_bytes();
Ok(ParquetPushDecoder {
state: ParquetDecoderState::ReadingRowGroup {
@@ -527,6 +535,15 @@ impl ParquetPushDecoder {
/// Note this can be the entire file or just a part of it. If it is part of the file,
/// the ranges should correspond to the data ranges requested by the decoder.
///
+ /// # Memory Usage
+ ///
+ /// The decoder releases the pushed bytes that it will not read again (for
+ /// example, the bytes of a row group when it is done with that row group).
+ /// However, it does not release *all* pushed bytes: it keeps the bytes of
+ /// row groups that it does not read (for example, row groups without
+ /// selected rows, or after the limit). Call [`Self::clear_all_ranges`] to
+ /// release them.
+ ///
/// See example in [`ParquetPushDecoderBuilder`]
pub fn push_range(&mut self, range: Range<u64>, data: Bytes) -> Result<(), ParquetError> {
self.push_ranges(vec![range], vec![data])
@@ -666,9 +683,9 @@ impl ParquetPushDecoder {
///
/// The decoder's buffered bytes are carried across the rebuild: bytes
/// already fetched for row groups the new configuration still reads are
- /// not re-requested. Bytes the new configuration no longer needs stay
- /// buffered until [`clear_all_ranges`](Self::clear_all_ranges) is called
- /// or the rebuilt decoder is dropped.
+ /// not re-requested. [`build`](ParquetPushDecoderBuilder::build) releases
+ /// the bytes that the new configuration does not read. See
+ /// [`push_range`](Self::push_range) for more details.
pub fn into_builder(self) -> Result<ParquetPushDecoderBuilder, ParquetError> {
self.state.into_builder()
}
@@ -974,6 +991,7 @@ mod test {
use super::*;
use crate::DecodeResult;
use crate::arrow::arrow_reader::{ArrowPredicateFn, RowFilter, RowSelection, RowSelector};
+ use crate::arrow::push_decoder::reader_builder::column_chunk_range;
use crate::arrow::push_decoder::{
ParquetPushDecoder, ParquetPushDecoderBuilder, RowGroupSelection,
};
@@ -1077,11 +1095,15 @@ mod test {
.unwrap();
assert_eq!(decoder.buffered_bytes(), test_file_len());
- // The current row group reader is built from the prefetched bytes, but
- // the speculative full-file range remains staged in the decoder.
+ // The current row group reader is built from the prefetched bytes.
+ // The decoder releases the bytes of the first row group, and the rest
+ // of the speculative full-file range remains staged in the decoder.
let batch1 = expect_data(decoder.try_decode());
assert_eq!(batch1, TEST_BATCH.slice(0, 100));
- assert_eq!(decoder.buffered_bytes(), test_file_len());
+ assert_eq!(
+ decoder.buffered_bytes(),
+ test_file_len() - test_file_row_group_bytes(0)
+ );
// All of the buffer is released
decoder.clear_all_ranges();
@@ -1219,6 +1241,115 @@ mod test {
assert_eq!(batch, TEST_BATCH.slice(100, 100));
}
+ /// Push one buffer for each `NeedsData`, from the start of the first
+ /// requested range to the end of the last one, and check that the
+ /// decoder releases the bytes of each row group when it is done with it.
+ ///
+ /// <https://github.com/apache/arrow-rs/issues/11233>
+ fn decode_with_coalesced_pushes(mut decoder: ParquetPushDecoder) -> Vec<RecordBatch> {
+ let mut batches = vec![];
+ loop {
+ match decoder.try_decode().unwrap() {
+ DecodeResult::NeedsData(ranges) => {
+ let start = ranges.iter().map(|r| r.start).min().unwrap();
+ let end = ranges.iter().map(|r| r.end).max().unwrap();
+ let coalesced = start..end;
+ push_ranges_to_decoder(&mut decoder, vec![coalesced]);
+ }
+ DecodeResult::Data(batch) => {
+ // The test file has one batch per row group.
+ assert_eq!(decoder.buffered_bytes(), 0);
+ batches.push(batch);
+ }
+ DecodeResult::Finished => return batches,
+ }
+ }
+ }
+
+ /// The requested ranges are adjacent, so the pushed buffer holds only
+ /// requested bytes.
+ #[test]
+ fn test_decoder_releases_coalesced_push() {
+ let decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
+ .unwrap()
+ .build()
+ .unwrap();
+ let batches = decode_with_coalesced_pushes(decoder);
+ let all_output = concat_batches(&TEST_BATCH.schema(), &batches).unwrap();
+ assert_eq!(all_output, *TEST_BATCH);
+ }
+
+ /// Column "b" is not read, so the pushed buffer also holds the bytes of
+ /// "b", which the decoder did not request.
+ #[test]
+ fn test_decoder_releases_unrequested_bytes_of_row_group() {
+ let builder =
+ ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap();
+ let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
+ let decoder = builder
+ .with_projection(ProjectionMask::columns(&schema_descr, ["a", "c"]))
+ .build()
+ .unwrap();
+ let batches = decode_with_coalesced_pushes(decoder);
+ let all_output = concat_batches(&batches[0].schema(), &batches).unwrap();
+ assert_eq!(all_output, TEST_BATCH.project(&[0, 2]).unwrap());
+ }
+
+ /// A plan that reads a row group two times keeps its bytes until the
+ /// second read.
+ #[test]
+ fn test_decoder_keeps_bytes_of_a_row_group_that_is_read_again() {
+ let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
+ .unwrap()
+ .with_row_groups(vec![0, 1, 0])
+ .build()
+ .unwrap();
+ prefetch_test_file(&mut decoder);
+ // Row group 0 is read again, so the decoder keeps its bytes.
+ assert_eq!(expect_data(decoder.try_decode()), TEST_BATCH.slice(0, 200));
+ assert_eq!(decoder.buffered_bytes(), test_file_len());
+
+ assert_eq!(
+ expect_data(decoder.try_decode()),
+ TEST_BATCH.slice(200, 200)
+ );
+ assert_eq!(
+ decoder.buffered_bytes(),
+ test_file_len() - test_file_row_group_bytes(1)
+ );
+
+ assert_eq!(expect_data(decoder.try_decode()), TEST_BATCH.slice(0, 200));
+ assert_eq!(
+ decoder.buffered_bytes(),
+ test_file_len() - test_file_row_group_bytes(0) - test_file_row_group_bytes(1)
+ );
+ expect_finished(decoder.try_decode());
+ }
+
+ /// With a predicate, the decoder requests the bytes of the predicate
+ /// column first, then the bytes of the other output columns.
+ #[test]
+ fn test_decoder_releases_coalesced_push_with_predicate() {
+ let builder =
+ ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap();
+ let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
+ let row_filter_a = ArrowPredicateFn::new(
+ ProjectionMask::columns(&schema_descr, ["a"]),
+ |batch: RecordBatch| {
+ let scalar_250 = Int64Array::new_scalar(250);
+ let column = batch.column(0).as_primitive::<Int64Type>();
+ gt(column, &scalar_250)
+ },
+ );
+ let decoder = builder
+ .with_row_filter(RowFilter::new(vec![Box::new(row_filter_a)]))
+ .build()
+ .unwrap();
+ let batches = decode_with_coalesced_pushes(decoder);
+ let all_output = concat_batches(&TEST_BATCH.schema(), &batches).unwrap();
+ assert_eq!(all_output, TEST_BATCH.slice(251, 149));
+ }
+
/// Decode multiple columns "a" and "b", expect that the decoder requests
/// only a single request per row group
#[test]
@@ -2644,9 +2775,11 @@ mod test {
let buffered = decoder.buffered_bytes();
assert!(buffered > 0);
- // Rebuilding via into_builder keeps the staged bytes.
+ // Rebuilding via into_builder keeps the staged bytes that RG1 reads,
+ // and releases the other bytes of the prefetched file.
let mut decoder = decoder.into_builder().unwrap().build().unwrap();
- assert_eq!(decoder.buffered_bytes(), buffered);
+ assert!(decoder.buffered_bytes() < buffered);
+ assert_eq!(decoder.buffered_bytes(), test_file_row_group_bytes(1));
// RG1's bytes are already buffered, so it decodes without a
// `NeedsData` round-trip.
@@ -2657,6 +2790,78 @@ mod test {
expect_finished(decoder.try_next_reader());
}
+ /// A rebuilt decoder releases the bytes of a row group that it skips,
+ /// and keeps the bytes of the columns that it reads.
+ #[test]
+ fn test_into_builder_releases_bytes_of_skipped_row_groups() {
+ let builder =
+ ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap();
+ let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
+ let mut decoder = builder
+ .with_row_groups(vec![1])
+ .with_projection(ProjectionMask::columns(&schema_descr, ["a", "b"]))
+ .build()
+ .unwrap();
+ prefetch_test_file(&mut decoder);
+ assert!(decoder.is_at_row_group_boundary());
+
+ // Skip RG1 and read RG0, with column "c" too.
+ let mut decoder = decoder
+ .into_builder()
+ .unwrap()
+ .with_row_groups(vec![0])
+ .with_projection(ProjectionMask::all())
+ .build()
+ .unwrap();
+ assert_eq!(decoder.buffered_bytes(), test_file_row_group_bytes(0));
+
+ let reader = expect_data(decoder.try_next_reader());
+ let batches: Vec<_> = reader.collect::<Result<_, _>>().unwrap();
+ let batch = concat_batches(&TEST_BATCH.schema(), &batches).unwrap();
+ assert_eq!(batch, TEST_BATCH.slice(0, 200));
+ expect_finished(decoder.try_next_reader());
+ }
+
+ /// A rebuilt decoder keeps the bytes of a column that only a predicate
+ /// reads.
+ #[test]
+ fn test_into_builder_keeps_bytes_of_predicate_columns() {
+ let builder =
+ ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap();
+ let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
+ let row_filter_a = ArrowPredicateFn::new(
+ ProjectionMask::columns(&schema_descr, ["a"]),
+ |batch: RecordBatch| {
+ let scalar_250 = Int64Array::new_scalar(250);
+ let column = batch.column(0).as_primitive::<Int64Type>();
+ gt(column, &scalar_250)
+ },
+ );
+ let mut decoder = builder
+ .with_row_groups(vec![1])
+ .with_projection(ProjectionMask::columns(&schema_descr, ["b"]))
+ .with_row_filter(RowFilter::new(vec![Box::new(row_filter_a)]))
+ .build()
+ .unwrap();
+ prefetch_test_file(&mut decoder);
+
+ // The rebuilt decoder keeps the bytes of "a" (the predicate) and "b"
+ // (the output) of RG1, and releases the bytes of "c".
+ let mut decoder = decoder.into_builder().unwrap().build().unwrap();
+ let column_bytes = |column_idx: usize| {
+ let range =
+ column_chunk_range(test_file_parquet_metadata().row_group(1).column(column_idx));
+ range.end - range.start
+ };
+ assert_eq!(decoder.buffered_bytes(), column_bytes(0) + column_bytes(1));
+
+ let reader = expect_data(decoder.try_next_reader());
+ let batches: Vec<_> = reader.collect::<Result<_, _>>().unwrap();
+ let batch = concat_batches(&batches[0].schema(), &batches).unwrap();
+ assert_eq!(batch, TEST_BATCH.slice(251, 149).project(&[1]).unwrap());
+ expect_finished(decoder.try_next_reader());
+ }
+
#[test]
fn test_into_builder_preserves_remaining_row_group_local_selections() {
let bitmap_selection = RowSelection::from_boolean_buffer(BooleanBuffer::from(
@@ -2934,6 +3139,20 @@ mod test {
decoder
}
+ /// The number of bytes in all column chunks of a row group of the test
+ /// file.
+ fn test_file_row_group_bytes(row_group_idx: usize) -> u64 {
+ test_file_parquet_metadata()
+ .row_group(row_group_idx)
+ .columns()
+ .iter()
+ .map(|column| {
+ let range = column_chunk_range(column);
+ range.end - range.start
+ })
+ .sum()
+ }
+
fn push_ranges_to_decoder(decoder: &mut ParquetPushDecoder, ranges: Vec<Range<u64>>) {
let data = ranges
.iter()
diff --git a/parquet/src/arrow/push_decoder/reader_builder/mod.rs b/parquet/src/arrow/push_decoder/reader_builder/mod.rs
index 8f122863be..52fa8320a1 100644
--- a/parquet/src/arrow/push_decoder/reader_builder/mod.rs
+++ b/parquet/src/arrow/push_decoder/reader_builder/mod.rs
@@ -33,8 +33,8 @@ use crate::arrow::push_decoder::reader_builder::filter::CacheInfo;
use crate::arrow::push_decoder::scan_plan::{BudgetedReadPlan, RowBudget};
use crate::arrow::schema::ParquetField;
use crate::errors::ParquetError;
-use crate::file::metadata::ParquetMetaData;
use crate::file::metadata::page_index::RowGroupPageIndex;
+use crate::file::metadata::{ColumnChunkMetaData, ParquetMetaData};
use crate::util::push_buffers::PushBuffers;
use bytes::Bytes;
use data::DataRequest;
@@ -91,6 +91,26 @@ enum RowGroupDecoderState {
Finished,
}
+impl RowGroupDecoderState {
+ /// The index of the row group, if one is active.
+ fn row_group_idx(&self) -> Option<usize> {
+ match self {
+ Self::Start { row_group_info }
+ | Self::Filters { row_group_info, .. }
+ | Self::WaitingOnFilterData { row_group_info, .. }
+ | Self::StartData { row_group_info, .. }
+ | Self::WaitingOnData { row_group_info, .. } => Some(row_group_info.row_group_idx),
+ Self::Finished => None,
+ }
+ }
+}
+
+/// The byte range of a column chunk in the file.
+pub(crate) fn column_chunk_range(column: &ColumnChunkMetaData) -> Range<u64> {
+ let (start, length) = column.byte_range();
+ start..start + length
+}
+
#[derive(Debug)]
pub(crate) enum RowGroupBuildResult {
/// The active row group is complete without producing a reader.
@@ -756,6 +776,58 @@ impl RowGroupReaderBuilder {
Ok(result)
}
+ /// The index of the active row group, if any.
+ pub(crate) fn active_row_group_idx(&self) -> Option<usize> {
+ self.state
+ .as_ref()
+ .and_then(RowGroupDecoderState::row_group_idx)
+ }
+
+ /// Remove the buffered bytes of all column chunks of a row group. This
+ /// includes bytes that the caller pushed but the decoder did not request,
+ /// for example the bytes between the requested ranges of a larger pushed
+ /// buffer.
+ pub(crate) fn release_row_group(&mut self, row_group_idx: usize) {
+ let ranges: Vec<Range<u64>> = self
+ .metadata
+ .row_group(row_group_idx)
+ .columns()
+ .iter()
+ .map(column_chunk_range)
+ .collect();
+ self.buffers.release_ranges(&ranges);
+ }
+
+ /// Remove the buffered bytes outside the read column chunks of
+ /// `row_groups`. A column chunk is read if the output or a predicate
+ /// reads its column. Indexes that are not in the file are ignored.
+ pub(crate) fn release_unread_bytes(&mut self, row_groups: impl IntoIterator<Item = usize>) {
+ if self.buffers.buffered_bytes() == 0 {
+ return;
+ }
+ let mut read_columns = self.projection.clone();
+ if let Some(filter) = &self.filter {
+ for predicate in &filter.predicates {
+ read_columns.union(predicate.projection());
+ }
+ }
+ let mut keep = vec![];
+ for row_group_idx in row_groups {
+ let Some(row_group) = self.metadata.row_groups().get(row_group_idx) else {
+ continue;
+ };
+ keep.extend(
+ row_group
+ .columns()
+ .iter()
+ .enumerate()
+ .filter(|(column_idx, _)| read_columns.leaf_included(*column_idx))
+ .map(|(_, column)| column_chunk_range(column)),
+ );
+ }
+ self.buffers.retain_ranges(&keep);
+ }
+
/// Which columns should be cached?
///
/// Returns the columns that are used by the filters *and* then used in the
diff --git a/parquet/src/arrow/push_decoder/remaining.rs b/parquet/src/arrow/push_decoder/remaining.rs
index 77f0e2f36f..afe69455c9 100644
--- a/parquet/src/arrow/push_decoder/remaining.rs
+++ b/parquet/src/arrow/push_decoder/remaining.rs
@@ -151,9 +151,9 @@ impl RemainingRowGroups {
/// when no row groups remain, or when every remaining row group
/// would be skipped under the current selection/budget.
///
- /// Cost: one clone of the queued row-group plan and selections per call
- /// (the frontier is cloned so the real advance logic can run
- /// non-destructively). For callers that peek once per row-group boundary
+ /// Cost: one clone of the queued row-group plan, selections and
+ /// row-group occurrence counts per call (the frontier is cloned so the
+ /// real advance logic can run non-destructively). For callers that peek once per row-group boundary
/// this is O(remaining row groups + selectors) per boundary.
pub fn peek_next_row_group(&self) -> Result<Option<usize>, ParquetError> {
if self.row_group_reader_builder.has_active_row_group() {
@@ -162,6 +162,25 @@ impl RemainingRowGroups {
self.frontier.peek_next_row_group()
}
+ /// Release the buffered bytes that are outside the read column chunks of
+ /// the queued row groups. The decoder does not read these bytes.
+ pub fn release_unread_bytes(&mut self) {
+ self.row_group_reader_builder
+ .release_unread_bytes(self.frontier.queued_row_groups());
+ }
+
+ /// Release the buffered bytes of a row group that is done, unless the
+ /// queue reads it again. The reader of the row group holds its own
+ /// copies of the bytes that it reads.
+ fn release_row_group(&mut self, row_group_idx: Option<usize>) {
+ if let Some(row_group_idx) = row_group_idx
+ && !self.frontier.is_queued(row_group_idx)
+ {
+ self.row_group_reader_builder
+ .release_row_group(row_group_idx);
+ }
+ }
+
/// returns [`ParquetRecordBatchReader`] suitable for reading the next
/// group of rows from the Parquet data, or the list of data ranges still
/// needed to proceed
@@ -191,10 +210,12 @@ impl RemainingRowGroups {
}
}
+ let row_group_idx = self.row_group_reader_builder.active_row_group_idx();
match self.row_group_reader_builder.try_build()? {
RowGroupBuildResult::Finished { remaining_budget } => {
self.frontier
.update_budget_after_row_group(remaining_budget);
+ self.release_row_group(row_group_idx);
// reader is done, proceed to the next row group
}
RowGroupBuildResult::NeedsData(ranges) => {
@@ -207,6 +228,7 @@ impl RemainingRowGroups {
} => {
self.frontier
.update_budget_after_row_group(remaining_budget);
+ self.release_row_group(row_group_idx);
// ready to read the row group
return Ok(DecodeResult::Data(batch_reader));
}
diff --git a/parquet/src/arrow/push_decoder/scan_plan/frontier.rs b/parquet/src/arrow/push_decoder/scan_plan/frontier.rs
index c35d91c1d9..7802922266 100644
--- a/parquet/src/arrow/push_decoder/scan_plan/frontier.rs
+++ b/parquet/src/arrow/push_decoder/scan_plan/frontier.rs
@@ -20,7 +20,7 @@
use crate::arrow::arrow_reader::{RowGroupPlan, RowGroupSelection, RowSelection};
use crate::errors::ParquetError;
use crate::file::metadata::ParquetMetaData;
-use std::collections::VecDeque;
+use std::collections::{HashMap, VecDeque};
use std::sync::Arc;
use super::budget::RowBudget;
@@ -48,7 +48,79 @@ pub(crate) struct NextRowGroup {
/// Row groups and selections that have not yet been handed to the row-group
/// reader builder.
#[derive(Debug, Clone)]
-enum QueuedRowGroups {
+struct QueuedRowGroups {
+ queue: Queue,
+ /// The number of times each row group is in `queue`, so that
+ /// [`Self::contains`] does not scan the queue.
+ counts: HashMap<usize, usize>,
+}
+
+impl QueuedRowGroups {
+ /// Validate and queue a row-group plan for `parquet_metadata`.
+ fn try_new(
+ parquet_metadata: &ParquetMetaData,
+ row_group_plan: RowGroupPlan,
+ ) -> Result<Self, ParquetError> {
+ let queue = Queue::try_new(parquet_metadata, row_group_plan)?;
+ let mut counts = HashMap::new();
+ for row_group_idx in queue.row_group_indices() {
+ *counts.entry(row_group_idx).or_default() += 1;
+ }
+ Ok(Self { queue, counts })
+ }
+
+ /// Convert the remaining queue back into a builder configuration.
+ fn into_plan(self) -> RowGroupPlan {
+ self.queue.into_plan()
+ }
+
+ fn front(&self) -> Option<usize> {
+ self.queue.front()
+ }
+
+ /// The row groups in the queue, in no specific order and without
+ /// duplicates.
+ fn row_groups(&self) -> impl Iterator<Item = usize> + '_ {
+ self.counts.keys().copied()
+ }
+
+ /// Returns `true` if `row_group_idx` is in the queue.
+ fn contains(&self, row_group_idx: usize) -> bool {
+ self.counts.contains_key(&row_group_idx)
+ }
+
+ fn len(&self) -> usize {
+ self.queue.len()
+ }
+
+ fn clear(&mut self) {
+ self.queue.clear();
+ self.counts.clear();
+ }
+
+ /// See [`Queue::global_selection_is_exhausted`].
+ fn global_selection_is_exhausted(&self) -> bool {
+ self.queue.global_selection_is_exhausted()
+ }
+
+ /// Remove the front row group and return its local selection.
+ fn pop_front_selection(&mut self, row_count: usize) -> Option<RowSelection> {
+ if let Some(row_group_idx) = self.queue.front() {
+ match self.counts.get_mut(&row_group_idx) {
+ Some(1) => {
+ self.counts.remove(&row_group_idx);
+ }
+ Some(count) => *count -= 1,
+ None => debug_assert!(false, "row group {row_group_idx} is queued but not counted"),
+ }
+ }
+ self.queue.pop_front_selection(row_count)
+ }
+}
+
+/// The queue of [`QueuedRowGroups`].
+#[derive(Debug, Clone)]
+enum Queue {
/// One selection cursor spans all queued row groups.
Global {
row_groups: VecDeque<usize>,
@@ -58,7 +130,7 @@ enum QueuedRowGroups {
PerRowGroup(VecDeque<RowGroupSelection>),
}
-impl QueuedRowGroups {
+impl Queue {
/// Validate and queue a row-group plan for `parquet_metadata`.
fn try_new(
parquet_metadata: &ParquetMetaData,
@@ -117,6 +189,17 @@ impl QueuedRowGroups {
}
}
+ /// The queued row group indexes, in order.
+ fn row_group_indices(&self) -> Vec<usize> {
+ match self {
+ Self::Global { row_groups, .. } => row_groups.iter().copied().collect(),
+ Self::PerRowGroup(row_groups) => row_groups
+ .iter()
+ .map(|row_group| row_group.row_group_index)
+ .collect(),
+ }
+ }
+
fn len(&self) -> usize {
match self {
Self::Global { row_groups, .. } => row_groups.len(),
@@ -202,6 +285,16 @@ impl RowGroupFrontier {
})
}
+ /// The queued row groups, in no specific order and without duplicates.
+ pub(crate) fn queued_row_groups(&self) -> impl Iterator<Item = usize> + '_ {
+ self.queued.row_groups()
+ }
+
+ /// Returns `true` if `row_group_idx` is still queued to be read.
+ pub(crate) fn is_queued(&self, row_group_idx: usize) -> bool {
+ self.queued.contains(row_group_idx)
+ }
+
pub(crate) fn update_budget_after_row_group(&mut self, budget: RowBudget) {
self.budget = budget;
}
@@ -456,6 +549,89 @@ mod tests {
assert_eq!(exhausted_budget.queued.len(), 0);
}
+ #[test]
+ fn queued_row_groups_count_repeated_row_groups() {
+ let metadata = test_file_parquet_metadata();
+ let budget = RowBudget::new(None, None);
+
+ // Row group 1 is skipped the first time. A row group stays queued
+ // until its last occurrence is popped.
+ let mut local = RowGroupFrontier::new(
+ Arc::clone(&metadata),
+ RowGroupPlan::PerRowGroup(vec![
+ RowGroupSelection::new(0, None),
+ RowGroupSelection::new(1, Some(RowSelection::from(vec![RowSelector::skip(200)]))),
+ RowGroupSelection::new(0, None),
+ RowGroupSelection::new(1, None),
+ ]),
+ budget,
+ false,
+ )
+ .unwrap();
+ let queued = |frontier: &RowGroupFrontier| {
+ (frontier.queued.contains(0), frontier.queued.contains(1))
+ };
+ assert_eq!(queued(&local), (true, true));
+
+ assert_eq!(
+ local
+ .next_readable_row_group()
+ .unwrap()
+ .unwrap()
+ .row_group_idx,
+ 0
+ );
+ assert_eq!(queued(&local), (true, true));
+ // Peeking does not change the counts of the frontier.
+ assert_eq!(local.peek_next_row_group().unwrap(), Some(0));
+ assert_eq!(queued(&local), (true, true));
+
+ // Skips the first occurrence of row group 1.
+ assert_eq!(
+ local
+ .next_readable_row_group()
+ .unwrap()
+ .unwrap()
+ .row_group_idx,
+ 0
+ );
+ assert_eq!(queued(&local), (false, true));
+
+ assert_eq!(
+ local
+ .next_readable_row_group()
+ .unwrap()
+ .unwrap()
+ .row_group_idx,
+ 1
+ );
+ assert_eq!(queued(&local), (false, false));
+ assert!(local.queued.counts.is_empty());
+ assert!(local.next_readable_row_group().unwrap().is_none());
+
+ // An exhausted limit clears the queue and its counts.
+ let mut global = RowGroupFrontier::new(
+ metadata,
+ global_plan(Some(vec![0, 1, 0]), None),
+ RowBudget::new(None, Some(200)),
+ false,
+ )
+ .unwrap();
+ assert_eq!(
+ global
+ .next_readable_row_group()
+ .unwrap()
+ .unwrap()
+ .row_group_idx,
+ 0
+ );
+ assert_eq!(queued(&global), (true, true));
+ global.update_budget_after_row_group(RowBudget::new(None, Some(0)));
+ assert!(global.next_readable_row_group().unwrap().is_none());
+ assert_eq!(queued(&global), (false, false));
+ assert!(global.queued.counts.is_empty());
+ }
+
#[test]
fn frontier_reports_invalid_global_row_group_while_peeking() {
let metadata = test_file_parquet_metadata();
diff --git a/parquet/src/util/push_buffers.rs b/parquet/src/util/push_buffers.rs
index d1a4c250e2..a7c1675335 100644
--- a/parquet/src/util/push_buffers.rs
+++ b/parquet/src/util/push_buffers.rs
@@ -66,6 +66,9 @@ pub struct PushBuffers {
/// when no earlier buffer can reach the requested range, so it checks one
/// or two buffers, not all of them. See [`Self::find`].
max_len: u64,
+ /// The sum of the lengths of `ranges`, kept up to date so that
+ /// [`Self::buffered_bytes`] does not scan all ranges.
+ buffered_bytes: u64,
}
impl Display for PushBuffers {
@@ -105,6 +108,7 @@ impl PushBuffers {
ranges: Vec::new(),
buffers: Vec::new(),
max_len: 0,
+ buffered_bytes: 0,
}
}
@@ -152,6 +156,7 @@ impl PushBuffers {
// Thus, ranges pushed in file order go at the end.
let idx = self.ranges.partition_point(|r| r.start <= range.start);
self.max_len = self.max_len.max(expected);
+ self.buffered_bytes += expected;
self.ranges.insert(idx, range);
self.buffers.insert(idx, buffer);
Ok(())
@@ -212,10 +217,15 @@ impl PushBuffers {
self
}
- /// Return the total of all buffered ranges
+ /// Return the total length of all buffered ranges.
+ ///
+ /// This counts the bytes that the decoder can still read, not the
+ /// allocated memory. [`Self::release_ranges`] keeps the remaining parts of
+ /// a buffer as slices of the pushed [`Bytes`], and the allocation of a
+ /// pushed [`Bytes`] is freed only after all of its parts are released.
#[cfg(feature = "arrow")]
pub(crate) fn buffered_bytes(&self) -> u64 {
- self.ranges.iter().map(|r| r.end - r.start).sum()
+ self.buffered_bytes
}
/// Clear any range and corresponding buffer that is exactly in the ranges_to_clear
@@ -237,6 +247,11 @@ impl PushBuffers {
}
self.ranges = ranges;
self.buffers = buffers;
+ self.buffered_bytes = self
+ .ranges
+ .iter()
+ .map(|r| r.end.saturating_sub(r.start))
+ .sum();
self.update_max_len();
}
@@ -250,16 +265,140 @@ impl PushBuffers {
self.max_len = self
.ranges
.iter()
- .map(|r| r.end - r.start)
+ .map(|r| r.end.saturating_sub(r.start))
.max()
.unwrap_or(0);
}
+ /// Remove all buffered bytes in `ranges`.
+ ///
+ /// A buffer that overlaps a range is trimmed or split, and the parts
+ /// outside the range are kept. The kept parts are zero-copy slices, which
+ /// record the bytes that the decoder can still read. Thus, the allocator
+ /// frees the memory of a pushed [`Bytes`] only after all of its parts are
+ /// removed.
+ ///
+ /// The buffers stay sorted by start, which lookups require (see
+ /// [`Self::find`]). If no buffer overlaps `ranges`, the buffers do not
+ /// change.
+ #[cfg(feature = "arrow")]
+ pub(crate) fn release_ranges(&mut self, ranges: &[Range<u64>]) {
+ let release = merge_ranges(ranges);
+ if release.is_empty() {
+ return;
+ }
+ // Most calls release ranges that are no longer buffered. Return
+ // before reallocating the buffers.
+ let overlaps = |range: &Range<u64>| {
+ !range.is_empty() && {
+ let first = release.partition_point(|r| r.end <= range.start);
+ release.get(first).is_some_and(|r| r.start < range.end)
+ }
+ };
+ if !self.ranges.iter().any(overlaps) {
+ return;
+ }
+ // Trim the buffers that overlap a released range in place, and do not
+ // clone or move the others. The parts after the first part of a split
+ // buffer go to `split`. A buffer that is released entirely becomes
+ // empty, and is removed below.
+ let mut split = vec![];
+ let mut emptied = false;
+ for (range, buffer) in self.ranges.iter_mut().zip(self.buffers.iter_mut()) {
+ if range.is_empty() {
+ emptied = true;
+ continue;
+ }
+ if !overlaps(range) {
+ continue;
+ }
+ let whole = range.clone();
+ let offset = |pos: u64| (pos - whole.start) as usize;
+ // The parts of `whole` between the released ranges.
+ let first = release.partition_point(|r| r.end <= whole.start);
+ let mut kept = 0;
+ let mut first_part = None;
+ let mut keep = |part: Range<u64>| {
+ kept += part.end - part.start;
+ if first_part.is_none() {
+ first_part = Some(part);
+ } else {
+ let data = buffer.slice(offset(part.start)..offset(part.end));
+ split.push((part, data));
+ }
+ };
+ let mut start = whole.start;
+ for r in release[first..].iter().take_while(|r| r.start < whole.end) {
+ if start < r.start {
+ keep(start..r.start);
+ }
+ start = r.end;
+ }
+ if start < whole.end {
+ keep(start..whole.end);
+ }
+ self.buffered_bytes -= (whole.end - whole.start) - kept;
+ match first_part {
+ Some(part) => {
+ buffer.truncate(offset(part.end));
+ let _ = buffer.split_to(offset(part.start));
+ *range = part;
+ }
+ None => {
+ *buffer = Bytes::new();
+ *range = whole.start..whole.start;
+ emptied = true;
+ }
+ }
+ }
+ if emptied {
+ // Remove the released buffers. As before, this also removes empty
+ // buffers that were pushed. `retain` visits each element exactly
+ // once in the original order, so both calls keep the same indices.
+ let mut keep = self.ranges.iter().map(|range| !range.is_empty());
+ self.buffers.retain(|_| keep.next().unwrap());
+ self.ranges.retain(|range| !range.is_empty());
+ }
+ for (range, buffer) in split {
+ self.ranges.push(range);
+ self.buffers.push(buffer);
+ }
+ // Split parts are appended, and if buffers overlap, the tail of a
+ // trimmed buffer can start after the start of the next buffer. The
+ // sort is stable.
+ if !self.ranges.is_sorted_by_key(|range| range.start) {
+ let mut parts: Vec<_> = std::mem::take(&mut self.ranges)
+ .into_iter()
+ .zip(std::mem::take(&mut self.buffers))
+ .collect();
+ parts.sort_by_key(|(range, _)| range.start);
+ (self.ranges, self.buffers) = parts.into_iter().unzip();
+ }
+ self.update_max_len();
+ }
+
+ /// Remove all buffered bytes outside `keep`.
+ ///
+ /// This calls [`Self::release_ranges`] with the complement of `keep`, so
+ /// the same rules apply to the kept parts.
+ #[cfg(feature = "arrow")]
+ pub(crate) fn retain_ranges(&mut self, keep: &[Range<u64>]) {
+ let mut release = vec![];
+ let mut start = 0;
+ for range in merge_ranges(keep) {
+ release.push(start..range.start);
+ start = range.end;
+ }
+ release.push(start..u64::MAX);
+ self.release_ranges(&release);
+ }
+
/// Clear all buffered ranges and their corresponding data
pub(crate) fn clear_all_ranges(&mut self) {
self.ranges.clear();
self.buffers.clear();
self.max_len = 0;
+ self.buffered_bytes = 0;
}
/// Panics if `ranges`, `buffers` and `max_len` do not agree, or if the
@@ -274,13 +413,33 @@ impl PushBuffers {
self.ranges
);
for (range, buffer) in self.ranges.iter().zip(&self.buffers) {
- assert_eq!(range.end - range.start, buffer.len() as u64);
+ assert_eq!(range.end.saturating_sub(range.start), buffer.len() as u64);
}
- let max_len = self.ranges.iter().map(|r| r.end - r.start).max();
+ let max_len = self
+ .ranges
+ .iter()
+ .map(|r| r.end.saturating_sub(r.start))
+ .max();
assert_eq!(self.max_len, max_len.unwrap_or(0));
}
}
+/// Sort `ranges` by start, remove empty ranges, and merge ranges that
+/// overlap or are adjacent.
+#[cfg(feature = "arrow")]
+fn merge_ranges(ranges: &[Range<u64>]) -> Vec<Range<u64>> {
+ let mut sorted: Vec<Range<u64>> = ranges.iter().filter(|r| !r.is_empty()).cloned().collect();
+ sorted.sort_unstable_by_key(|r| r.start);
+ let mut merged: Vec<Range<u64>> = Vec::with_capacity(sorted.len());
+ for range in sorted {
+ match merged.last_mut() {
+ Some(last) if range.start <= last.end => last.end = last.end.max(range.end),
+ _ => merged.push(range),
+ }
+ }
+ merged
+}
+
impl Length for PushBuffers {
fn len(&self) -> u64 {
self.file_len
@@ -334,6 +493,143 @@ impl ChunkReader for PushBuffers {
mod tests {
use super::*;
+ /// Release one range.
+ #[cfg(feature = "arrow")]
+ fn release(buffers: &mut PushBuffers, range: Range<u64>) {
+ buffers.release_ranges(std::slice::from_ref(&range));
+ buffers.assert_invariants();
+ }
+
+ /// The buffered ranges and their data.
+ #[cfg(feature = "arrow")]
+ fn buffered_parts(buffers: &PushBuffers) -> Vec<(Range<u64>, Bytes)> {
+ let ranges = buffers.ranges.iter().cloned();
+ ranges.zip(buffers.buffers.iter().cloned()).collect()
+ }
+
+ #[test]
+ #[cfg(feature = "arrow")]
+ fn release_ranges_without_overlap_keeps_buffers() {
+ let mut buffers = PushBuffers::new(100);
+ buffers
+ .push_range(20..24, Bytes::from_static(b"abcd"))
+ .unwrap();
+ buffers
+ .push_range(0..10, Bytes::from_static(b"0123456789"))
+ .unwrap();
+
+ release(&mut buffers, 90..95);
+ release(&mut buffers, 10..20);
+ assert_eq!(buffers.ranges, vec![0..10, 20..24]);
+ assert_eq!(buffers.buffered_bytes(), 14);
+
+ release(&mut buffers, 22..24);
+ assert_eq!(buffers.ranges, vec![0..10, 20..22]);
+ assert_eq!(buffers.buffered_bytes(), 12);
+
+ buffers.clear_all_ranges();
+ assert_eq!(buffers.buffered_bytes(), 0);
+ }
+
+ #[test]
+ #[cfg(feature = "arrow")]
+ fn release_ranges_trims_and_splits_buffers() {
+ let mut buffers = PushBuffers::new(100);
+ buffers
+ .push_range(0..10, Bytes::from_static(b"0123456789"))
+ .unwrap();
+ buffers
+ .push_range(20..24, Bytes::from_static(b"abcd"))
+ .unwrap();
+
+ // Split the first buffer, and do not change the second buffer.
+ release(&mut buffers, 3..5);
+ assert_eq!(buffers.buffered_bytes(), 12);
+ assert!(buffers.has_range(&(0..3)));
+ assert!(buffers.has_range(&(5..10)));
+ assert!(!buffers.has_range(&(3..4)));
+ assert_eq!(
+ buffers.get_bytes(5, 5).unwrap(),
+ Bytes::from_static(b"56789")
+ );
+ assert!(buffers.has_range(&(20..24)));
+
+ // One range that overlaps two buffers trims both of them.
+ release(&mut buffers, 8..22);
+ assert_eq!(buffers.buffered_bytes(), 3 + 3 + 2);
+ assert_eq!(buffers.get_bytes(5, 3).unwrap(), Bytes::from_static(b"567"));
+ assert_eq!(buffers.get_bytes(22, 2).unwrap(), Bytes::from_static(b"cd"));
+
+ // Ranges in any order, adjacent or empty, and bytes that are not
+ // buffered.
+ buffers.release_ranges(&[50..60, 1..2, 0..1, 7..7]);
+ assert_eq!(buffers.buffered_bytes(), 1 + 3 + 2);
+ assert_eq!(buffers.get_bytes(2, 1).unwrap(), Bytes::from_static(b"2"));
+
+ release(&mut buffers, 0..100);
+ assert_eq!(buffers.buffered_bytes(), 0);
+ }
+
+ #[test]
+ #[cfg(feature = "arrow")]
+ fn release_ranges_removes_exact_and_overlapping_buffers() {
+ let mut buffers = PushBuffers::new(100);
+ for range in [10..20, 0..30, 10..15, 40..50] {
+ let data = Bytes::from(vec![0u8; (range.end - range.start) as usize]);
+ buffers.push_range(range, data).unwrap();
+ }
+ buffers.release_ranges(&[40..50, 10..15]);
+ buffers.assert_invariants();
+ assert_eq!(buffers.buffered_bytes(), 10 + 5 + 15);
+ assert!(buffers.has_range(&(0..10)));
+ assert!(buffers.has_range(&(15..30)));
+ assert!(!buffers.has_range(&(10..15)));
+ assert!(!buffers.has_range(&(40..50)));
+ }
+
+ #[test]
+ #[cfg(feature = "arrow")]
+ fn retain_ranges_keeps_only_the_given_bytes() {
+ let mut buffers = PushBuffers::new(100);
+ buffers
+ .push_range(0..10, Bytes::from_static(b"0123456789"))
+ .unwrap();
+ buffers
+ .push_range(20..24, Bytes::from_static(b"abcd"))
+ .unwrap();
+ buffers
+ .push_range(30..32, Bytes::from_static(b"xy"))
+ .unwrap();
+ // Ranges in any order, overlapping, and outside the buffers.
+ buffers.retain_ranges(&[22..40, 2..4, 3..5, 8..9]);
+ assert_eq!(buffers.buffered_bytes(), 3 + 1 + 2 + 2);
+ assert_eq!(buffers.get_bytes(2, 3).unwrap(), Bytes::from_static(b"234"));
+ assert_eq!(buffers.get_bytes(8, 1).unwrap(), Bytes::from_static(b"8"));
+ assert!(!buffers.has_range(&(5..6)));
+ assert_eq!(buffers.get_bytes(22, 2).unwrap(), Bytes::from_static(b"cd"));
+ assert_eq!(buffers.get_bytes(30, 2).unwrap(), Bytes::from_static(b"xy"));
+ buffers.retain_ranges(&[]);
+ assert_eq!(buffers.buffered_bytes(), 0);
+ }
+
+ /// Lookups stay correct when the kept parts of overlapping buffers
+ /// interleave.
+ #[test]
+ #[cfg(feature = "arrow")]
+ fn retain_ranges_with_overlapping_buffers() {
+ let mut buffers = PushBuffers::new(100);
+ for range in [0..100, 3..70] {
+ let data = Bytes::from(vec![0u8; (range.end - range.start) as usize]);
+ buffers.push_range(range, data).unwrap();
+ }
+ buffers.retain_ranges(&[0..5, 50..60]);
+ buffers.assert_invariants();
+ assert_eq!(buffers.buffered_bytes(), 5 + 2 + 10 + 10);
+ assert!(buffers.has_range(&(0..5)));
+ assert!(buffers.has_range(&(50..60)));
+ assert!(!buffers.has_range(&(5..50)));
+ }
+
#[test]
fn push_range_accepts_matching_length() {
let mut buffers = PushBuffers::new(100);
@@ -487,4 +783,131 @@ mod tests {
"Parquet error: Number of ranges (2) must match number of buffers (1)"
);
}
+
+ /// The result of `release_ranges`, computed by rebuilding all buffers.
+ #[cfg(feature = "arrow")]
+ fn release_by_rebuild(
+ buffers: &PushBuffers,
+ release: &[Range<u64>],
+ ) -> Vec<(Range<u64>, Bytes)> {
+ let release = merge_ranges(release);
+ let mut parts = vec![];
+ for (range, buffer) in &buffered_parts(buffers) {
+ let mut start = range.start;
+ let mut keep = |part: Range<u64>| {
+ let offset = |pos: u64| (pos - range.start) as usize;
+ parts.push((
+ part.clone(),
+ buffer.slice(offset(part.start)..offset(part.end)),
+ ));
+ };
+ for r in release
+ .iter()
+ .filter(|r| r.start < range.end && r.end > range.start)
+ {
+ if start < r.start {
+ keep(start..r.start);
+ }
+ start = start.max(r.end);
+ }
+ if start < range.end {
+ keep(start..range.end);
+ }
+ }
+ parts.sort_by_key(|(range, _)| range.start);
+ parts
+ }
+
+ #[test]
+ #[cfg(feature = "arrow")]
+ #[expect(clippy::reversed_empty_ranges)]
+ fn release_ranges_splits_and_removes_empty_buffers() {
+ let mut buffers = PushBuffers::new(100);
+ buffers
+ .push_range(0..20, Bytes::from_static(b"abcdefghijklmnopqrst"))
+ .unwrap();
+ // An empty and an inverted range, which `push_range` accepts.
+ buffers.push_range(30..30, Bytes::new()).unwrap();
+ buffers.push_range(10..5, Bytes::new()).unwrap();
+
+ // Split one buffer into three parts.
+ buffers.release_ranges(&[3..5, 8..10]);
+ assert_eq!(buffers.ranges, vec![0..3, 5..8, 10..20]);
+ assert_eq!(buffers.buffered_bytes(), 16);
+ assert_eq!(buffers.get_bytes(5, 3).unwrap(), Bytes::from_static(b"fgh"));
+
+ // Release a whole buffer.
+ release(&mut buffers, 5..8);
+ assert_eq!(buffers.ranges, vec![0..3, 10..20]);
+ assert_eq!(buffers.buffered_bytes(), 13);
+ }
+
+ #[test]
+ #[cfg(feature = "arrow")]
+ fn fuzz_release_ranges_matches_rebuild() {
+ use rand::{RngExt, SeedableRng, rngs::StdRng};
+ let data: Vec<u8> = (0..=255).collect();
+ let mut rng = StdRng::seed_from_u64(42);
+ for _ in 0..2000 {
+ let mut buffers = PushBuffers::new(256);
+ for _ in 0..rng.random_range(0..6) {
+ let start = rng.random_range(0..100u64);
+ if rng.random_bool(0.1) {
+ // An inverted range, which `push_range` accepts.
+ let end = start.saturating_sub(rng.random_range(1..5u64));
+ buffers.push_range(start..end, Bytes::new()).unwrap();
+ continue;
+ }
+ let end = start + rng.random_range(0..20u64);
+ let bytes = Bytes::copy_from_slice(&data[start as usize..end as usize]);
+ buffers.push_range(start..end, bytes).unwrap();
+ }
+ for _ in 0..3 {
+ let release: Vec<_> = (0..rng.random_range(0..4))
+ .map(|_| {
+ let start = rng.random_range(0..110u64);
+ start..start + rng.random_range(0..15u64)
+ })
+ .collect();
+ let overlaps = buffers.ranges.iter().any(|range| {
+ !range.is_empty()
+ && release
+ .iter()
+ .any(|r| !r.is_empty() && r.start < range.end && r.end > range.start)
+ });
+ let before = buffered_parts(&buffers);
+ let expected = release_by_rebuild(&buffers, &release);
+
+ buffers.release_ranges(&release);
+ buffers.assert_invariants();
+ let actual = buffered_parts(&buffers);
+ if overlaps {
+ // The order of buffers with the same start is not
+ // specified.
+ let by_range = |mut parts: Vec<(Range<u64>, Bytes)>| {
+ parts.sort_by_key(|(range, _)| (range.start, range.end));
+ parts
+ };
+ assert_eq!(
+ by_range(actual.clone()),
+ by_range(expected),
+ "release {release:?} of {before:?}"
+ );
+ } else {
+ assert_eq!(actual, before, "release {release:?} of {before:?}");
+ }
+ for (range, buffer) in actual.iter().filter(|(range, _)| !range.is_empty()) {
+ assert_eq!(
+ &data[range.start as usize..range.end as usize],
+ buffer.as_ref()
+ );
+ }
+ let total: u64 = actual
+ .iter()
+ .map(|(r, _)| r.end.saturating_sub(r.start))
+ .sum();
+ assert_eq!(buffers.buffered_bytes(), total);
+ }
+ }
+ }
}