adriangb commented on code in PR #11237:
URL: https://github.com/apache/arrow-rs/pull/11237#discussion_r4211041578
##########
parquet/src/util/push_buffers.rs:
##########
@@ -275,6 +320,149 @@ mod tests {
assert!(!buffers.has_range(&(10..20)));
}
+ /// The bytes of a fake file: byte `i` is `i % 251`, so any slice of it is
+ /// easy to build and to check.
+ fn file_bytes(range: Range<u64>) -> Bytes {
+ range.map(|i| (i % 251) as u8).collect::<Vec<u8>>().into()
+ }
+
+ fn push(buffers: &mut PushBuffers, range: Range<u64>) {
+ buffers
+ .push_range(range.clone(), file_bytes(range))
+ .unwrap();
+ }
+
+ #[track_caller]
+ fn assert_invariants(buffers: &PushBuffers) {
Review Comment:
Resolved in
https://github.com/apache/arrow-rs/pull/11237/commits/53d9331fb2414016e0e8320da74cca84869f7eb7
##########
parquet/src/util/push_buffers.rs:
##########
@@ -132,20 +147,44 @@ impl PushBuffers {
range.end
));
}
- self.ranges.push(range);
- self.buffers.push(buffer);
+ // Insert after all buffers that start at or before `range.start`.
+ // 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.ranges.insert(idx, range);
+ self.buffers.insert(idx, buffer);
Ok(())
}
/// Returns true if the Buffers contains data for the given range
pub(crate) fn has_range(&self, range: &Range<u64>) -> bool {
- self.ranges
- .iter()
- .any(|r| r.start <= range.start && r.end >= range.end)
+ self.find(range.start, range.end).is_some()
}
- fn iter(&self) -> impl Iterator<Item = (&Range<u64>, &Bytes)> {
- self.ranges.iter().zip(self.buffers.iter())
+ /// Returns the index of a buffer that contains all bytes of `start..end`,
+ /// if any.
+ fn find(&self, start: u64, end: u64) -> Option<usize> {
+ // Ranges can overlap. Thus, the buffer that starts nearest to `start`
+ // is not always the buffer that contains `start..end`:
+ //
+ // buffers: 0..100 ├────────────────────────────────┤
+ // 50..60 ├───┤
+ // 55..58 ├┤
+ // find: 55..90 ├──────────┤ only 0..100
contains it
+ //
+ // Scan back from the last buffer that starts at or before `start`.
+ // Stop at the first buffer that starts more than `max_len` bytes
+ // before `end`: it and all buffers before it end before `end`. The
Review Comment:
Resolved in
https://github.com/apache/arrow-rs/pull/11237/commits/f98e54729620759ceeaa5da2bdb3e8293c0736ee
##########
parquet/src/util/push_buffers.rs:
##########
@@ -41,16 +41,30 @@ use std::ops::Range;
///
/// Thus, the implementation defers to the caller to coalesce subsequent
requests
/// if desired.
+///
+/// # Ordering
+///
+/// The buffers are sorted by the start of their range. Thus, lookups use a
+/// binary search, not a scan of all buffers. Pushed ranges can overlap.
#[derive(Debug, Clone, Default)]
pub struct PushBuffers {
/// the virtual "offset" of this buffers (added to any request)
offset: u64,
/// The total length of the file being decoded
file_len: u64,
- /// The ranges of data that are available for decoding (not adjusted for
offset)
+ /// The ranges of data that are available for decoding (not adjusted for
+ /// offset), sorted by `start`
ranges: Vec<Range<u64>>,
- /// The buffers of data that can be used to decode the Parquet file
+ /// The buffers of data that can be used to decode the Parquet file, in the
+ /// same order as `ranges`
buffers: Vec<Bytes>,
+ /// The maximum length of the ranges in `ranges`. A buffer that starts
Review Comment:
Resolved in
https://github.com/apache/arrow-rs/pull/11237/commits/ae56305f46a6baf5a03bd41be443c8d5fe4749e4
##########
parquet/src/util/push_buffers.rs:
##########
@@ -275,6 +320,149 @@ mod tests {
assert!(!buffers.has_range(&(10..20)));
}
+ /// The bytes of a fake file: byte `i` is `i % 251`, so any slice of it is
+ /// easy to build and to check.
+ fn file_bytes(range: Range<u64>) -> Bytes {
+ range.map(|i| (i % 251) as u8).collect::<Vec<u8>>().into()
+ }
+
+ fn push(buffers: &mut PushBuffers, range: Range<u64>) {
+ buffers
+ .push_range(range.clone(), file_bytes(range))
+ .unwrap();
+ }
+
+ #[track_caller]
+ fn assert_invariants(buffers: &PushBuffers) {
+ assert_eq!(buffers.ranges.len(), buffers.buffers.len());
+ assert!(
+ buffers.ranges.is_sorted_by_key(|r| r.start),
+ "not sorted: {:?}",
+ buffers.ranges
+ );
+ for (range, buffer) in buffers.ranges.iter().zip(&buffers.buffers) {
+ assert_eq!(*buffer, file_bytes(range.clone()));
+ assert!(range.end - range.start <= buffers.max_len);
+ }
+ let max_len = buffers.ranges.iter().map(|r| r.end - r.start).max();
+ assert_eq!(buffers.max_len, max_len.unwrap_or(0));
+ }
+
+ #[test]
+ fn out_of_order_pushes_are_sorted() {
Review Comment:
Resolved in
https://github.com/apache/arrow-rs/pull/11237/commits/1fe89c90131d0247cbf044ec4d9d927456fede8a
##########
parquet/src/util/push_buffers.rs:
##########
@@ -275,6 +320,149 @@ mod tests {
assert!(!buffers.has_range(&(10..20)));
}
+ /// The bytes of a fake file: byte `i` is `i % 251`, so any slice of it is
+ /// easy to build and to check.
+ fn file_bytes(range: Range<u64>) -> Bytes {
+ range.map(|i| (i % 251) as u8).collect::<Vec<u8>>().into()
+ }
+
+ fn push(buffers: &mut PushBuffers, range: Range<u64>) {
+ buffers
+ .push_range(range.clone(), file_bytes(range))
+ .unwrap();
+ }
+
+ #[track_caller]
+ fn assert_invariants(buffers: &PushBuffers) {
+ assert_eq!(buffers.ranges.len(), buffers.buffers.len());
+ assert!(
+ buffers.ranges.is_sorted_by_key(|r| r.start),
+ "not sorted: {:?}",
+ buffers.ranges
+ );
+ for (range, buffer) in buffers.ranges.iter().zip(&buffers.buffers) {
+ assert_eq!(*buffer, file_bytes(range.clone()));
+ assert!(range.end - range.start <= buffers.max_len);
+ }
+ let max_len = buffers.ranges.iter().map(|r| r.end - r.start).max();
+ assert_eq!(buffers.max_len, max_len.unwrap_or(0));
+ }
+
+ #[test]
+ fn out_of_order_pushes_are_sorted() {
+ let mut buffers = PushBuffers::new(1000);
+ for range in [50..60, 10..20, 30..40, 0..5, 90..100, 10..15] {
+ push(&mut buffers, range);
+ }
+ assert_invariants(&buffers);
+ assert_eq!(
+ buffers.ranges,
+ vec![0..5, 10..20, 10..15, 30..40, 50..60, 90..100]
+ );
+ assert!(buffers.has_range(&(12..18)));
+ assert!(buffers.has_range(&(30..40)));
+ assert!(!buffers.has_range(&(18..22)));
+ assert!(!buffers.has_range(&(60..61)));
+ assert_eq!(buffers.get_bytes(52, 4).unwrap(), file_bytes(52..56));
+ assert!(matches!(
+ buffers.get_bytes(40, 20),
+ Err(ParquetError::NeedMoreDataRange(r)) if r == (40..60)
+ ));
+ }
+
+ #[test]
+ fn overlapping_pushes_find_the_containing_buffer() {
+ let mut buffers = PushBuffers::new(1000);
+ // Small buffers inside a large one start closer to most offsets.
+ for range in [0..100, 10..20, 50..60, 55..58] {
+ push(&mut buffers, range);
+ }
+ assert_invariants(&buffers);
+ // 55..90 starts in 50..60 and 55..58, but only 0..100 contains it.
+ assert_eq!(buffers.get_bytes(55, 35).unwrap(), file_bytes(55..90));
+ assert!(buffers.has_range(&(0..100)));
+ assert!(buffers.has_range(&(99..100)));
+ assert!(!buffers.has_range(&(99..101)));
+
+ // `Read` finds the same buffer.
+ let mut reader = buffers.get_read(56).unwrap();
+ let mut out = [0u8; 30];
+ std::io::Read::read_exact(&mut reader, &mut out).unwrap();
+ assert_eq!(&out[..], &file_bytes(56..86)[..]);
+ }
+
+ #[test]
+ #[cfg(feature = "arrow")]
+ fn clear_ranges_drops_exact_matches_only() {
+ let mut buffers = PushBuffers::new(1000);
+ for range in [10..20, 0..30, 10..15, 40..50] {
+ push(&mut buffers, range);
+ }
+ buffers.clear_ranges(&[40..50, 10..15, 5..30]);
+ assert_invariants(&buffers);
+ assert_eq!(buffers.ranges, vec![0..30, 10..20]);
+ }
+
+ /// Random pushes, clears and lookups, compared with a list that is
+ /// scanned in full for each lookup.
+ #[test]
+ #[cfg(feature = "arrow")]
+ fn matches_a_linear_scan() {
Review Comment:
Resolved in
https://github.com/apache/arrow-rs/pull/11237/commits/a515b5c24285562f77969beff370dac00a1a0ce5
--
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]
