alamb commented on code in PR #11235:
URL: https://github.com/apache/arrow-rs/pull/11235#discussion_r4211759395
##########
parquet/src/arrow/push_decoder/mod.rs:
##########
@@ -527,6 +535,21 @@ 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.
///
+ /// When the decoder is done with a row group that it reads, it releases
Review Comment:
Maybe a heading like `# Memory Usage` would help readers break up this text
##########
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:
Review Comment:
I think this makes sense -- to put the cleaning as part of the builder
##########
parquet/src/arrow/push_decoder/mod.rs:
##########
@@ -527,6 +535,21 @@ 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.
///
+ /// When the decoder is done with a row group that it reads, it releases
+ /// the pushed bytes of all column chunks of that row group, unless the
+ /// scan reads the row group again. This is also true for bytes that the
Review Comment:
I wonder if we can make this comment simpler by summarizing they key points
rather than lots of detail.
For example, something like
> The reader releases byte ranges will never be read in the future.
However, it may not release *all* the ranges (for example, if there are
predicates or a limit that allow ranges to be skipped
##########
parquet/src/arrow/push_decoder/mod.rs:
##########
@@ -527,6 +535,21 @@ 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.
///
+ /// When the decoder is done with a row group that it reads, it releases
+ /// the pushed bytes of all column chunks of that row group, unless the
Review Comment:
the " unless the scan reads the row group again" seems redundant as if the
reader would scan the row group again, it wouldn't be done wiht it 🤔
##########
parquet/src/arrow/push_decoder/mod.rs:
##########
@@ -1077,11 +1102,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!(
Review Comment:
This is a test of the new behavior 👍
##########
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
Review Comment:
what is the importance of having the buffers sorted by start here? It seems
irrelevant to the release_ranges functionality
##########
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:
Can this code be simplified once we merge the change to keep the ranges
sorted in
- https://github.com/apache/arrow-rs/pull/11237
?
##########
parquet/src/arrow/push_decoder/mod.rs:
##########
@@ -527,6 +535,21 @@ 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.
///
+ /// When the decoder is done with a row group that it reads, it releases
+ /// the pushed bytes of all column chunks of that row group, unless the
+ /// scan reads the row group again. This is also true for bytes that the
+ /// decoder did not request, and for bytes in a buffer that is larger than
+ /// a requested range. The decoder does not release 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.
+ ///
+ /// Releasing bytes decreases [`Self::buffered_bytes`], but the memory of
+ /// a pushed [`Bytes`] is freed only after all zero-copy slices of it are
Review Comment:
The detail about memory management / reference counting of `Bytes` I think
is an implementation detail of Bytes that is not specific to `push_ranges`
##########
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
Review Comment:
I think I understand after some more review -- the idea is that the slices
are basically the accounting of the ranges that might still be used.
Maybe we can adjust the documentation for buffered bytes to also explain this
##########
parquet/src/arrow/push_decoder/mod.rs:
##########
@@ -1219,6 +1248,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
Review Comment:
these are nice tests
##########
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
Review Comment:
I am not sure what "whatever the shape of the pushed buffers" means -- can
we please clarify it or remove this phrase?
##########
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
Review Comment:
What is the value of slicing `Buffer`s ? It seems like it just reduces the
accuracy of the memory accounting as the allocation isn't released until all
ranges are relased 🤔
##########
parquet/src/util/push_buffers.rs:
##########
@@ -253,6 +399,126 @@ 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));
+ }
+
+ #[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![20..24, 0..10]);
+ assert_eq!(buffers.buffered_bytes(), 14);
+
+ // A release that overlaps a buffer sorts the buffers.
+ release(&mut buffers, 22..24);
Review Comment:
why are we testing that internally the buffers are sorted? That seems like
an implementation detail and not something that is part of the contract (aka it
is fine if it changes not to be sorted in the future, isn't it?
##########
parquet/src/arrow/push_decoder/mod.rs:
##########
@@ -666,9 +689,10 @@ 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: the bytes outside
Review Comment:
I think this could be simplified again (and the outside colum chunks it
doesn't read is too much detail)
Maybe something like
> [`build`](ParquetPushDecoderBuilder::build) releases
> the bytes that the new configuration will not read. See [`push_ranges`]
for more detaiks
##########
parquet/src/util/push_buffers.rs:
##########
@@ -286,4 +552,137 @@ 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 buffers.iter() {
+ 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 release_ranges_matches_rebuild() {
Review Comment:
Maybe calling this something called `fuzz` would make it clear it is doing
fuzz testing
--
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]
