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 <[email protected]> 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 <[email protected]> --- 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); + } + } + } }
