adriangb commented on code in PR #11235:
URL: https://github.com/apache/arrow-rs/pull/11235#discussion_r4211028986
##########
parquet/src/arrow/push_decoder/remaining.rs:
##########
@@ -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);
Review Comment:
**Start here.** This is the fix: when a row group is done, release all of
its column chunks. Before, the decoder released a buffer only if its range was
equal to a requested range.
##########
parquet/src/arrow/push_decoder/remaining.rs:
##########
@@ -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>) {
Review Comment:
A row group can be queued more than one time, so its bytes are kept while it
is still queued.
##########
parquet/src/arrow/push_decoder/scan_plan/frontier.rs:
##########
@@ -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 {
Review Comment:
This diff looks larger than it is. The old `QueuedRowGroups` enum is renamed
to `Queue`, and the only addition to it is `row_group_indices`. The new struct
adds occurrence counts so that `is_queued` does not scan the queue (from the
review by @HippoBaro).
##########
parquet/src/util/push_buffers.rs:
##########
@@ -180,15 +185,156 @@ impl PushBuffers {
new_buffers.push(buffer.clone());
}
}
+ self.buffered_bytes = new_ranges.iter().map(|r| r.end - r.start).sum();
self.ranges = new_ranges;
self.buffers = new_buffers;
}
+ /// Remove all buffered bytes in `ranges`, whatever the shape of the
+ /// pushed buffers.
+ ///
+ /// 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. Thus,
+ /// the allocator frees the memory of a pushed [`Bytes`] only after all of
+ /// its parts are removed.
+ ///
+ /// If the buffers are sorted by start, they stay sorted. The order of
+ /// buffers with the same start is not specified. If no buffer overlaps
+ /// `ranges`, the buffers do not change.
+ #[cfg(feature = "arrow")]
+ pub(crate) fn release_ranges(&mut self, ranges: &[Range<u64>]) {
Review Comment:
The new primitive. It trims or splits pushed buffers into zero-copy slices.
`release_ranges_matches_rebuild` compares it with a simple rebuild of all
buffers.
##########
parquet/src/arrow/push_decoder/mod.rs:
##########
@@ -347,14 +351,18 @@ 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,
RowBudget::new(offset, limit),
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();
Review Comment:
The second change, and the one behavior change for callers. With a new
builder the buffers are empty, so this returns early. It only has an effect
with `with_buffers` or `into_builder`.
##########
parquet/src/arrow/push_decoder/reader_builder/mod.rs:
##########
@@ -755,6 +775,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>) {
Review Comment:
Kept bytes = the column chunks of the output projection and of the predicate
columns, for the queued row groups.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
