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-11459-8dff89dbfed32b15115854e816c8e83f69f08226
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git

commit 236fc5a8d5a7d8667cfc481d9d63b6adc3860a79
Author: Adrian Garcia Badaracco <[email protected]>
AuthorDate: Wed Oct 14 22:46:03 2026 +0000

    perf(parquet): make `PushBuffers` release and clear cost independent of the 
buffer count (#11459)
    
    # Which issue does this PR close?
    
    - Closes https://github.com/apache/arrow-rs/issues/11458.
    
    Follow-up to https://github.com/apache/arrow-rs/pull/11235, which is
    merged. This PR has four commits:
    
    1.
    
https://github.com/apache/arrow-rs/pull/11459/commits/8087cef7fc49141c1e73e5030b5f572af6eef29e
    (new `push_decoder/prefetch` benchmark)
    2.
    
https://github.com/apache/arrow-rs/pull/11459/commits/c0857e763a6bdd19549813e221d605a678dc771b
    (binary search in `release_ranges`)
    3.
    
https://github.com/apache/arrow-rs/pull/11459/commits/04faa64177840ae1b1766b5ff7e269bd927d021d
    (binary search in `clear_ranges`)
    4.
    
https://github.com/apache/arrow-rs/pull/11459/commits/8b98d709ae5ba9c813fe4ae25fdbf53476f24ef1
    (`VecDeque` and `max_len_count`)
    
    # Rationale for this change
    
    When a caller pushes many buffers before the decoder reads them (for
    example, read-ahead for many row groups), the push decoder pays
    `O(buffers)` for each row group. Thus, the full scan is quadratic.
    
    The new `push_decoder/prefetch` benchmark shows this. It pushes one
    buffer for each column chunk of the file, then builds all row group
    readers. Time for each iteration on the GKE runner:
    
    | Buffers | Before | After | Change |
    |--|--|--|--|
    | 1,000 | 1.12 ms | 0.60 ms | 1.9x faster |
    | 10,000 | 56.6 ms | 6.3 ms | 9x faster |
    | 30,000 | 478 ms | 19.2 ms | 25x faster |
    
    After this PR, the time is linear in the number of buffers.
    
    The cost: the `VecDeque` commit makes the existing `push_decoder/1buf`
    and `push_decoder/Nbuf` cases 1% to 6% slower, because an index into a
    `VecDeque` costs more than an index into a `Vec`. Those cases hold few
    buffers at a time, so they get no gain from this PR.
    
    # What changes are included in this PR?
    
    For each row group, the decoder calls `clear_ranges` and then
    `release_ranges`. Each call visited all buffers:
    
    ```diff
     for each row group
       clear_ranges(requested ranges)
    -    rebuild all buffers, keep those that are not requested
    -    scan all buffers for max_len
    +    binary search for the buffers that start at each requested range
    +    remove them, move no other buffer (VecDeque)
    +    scan for max_len only if no buffer of that length remains
       release_ranges(column chunk ranges)
    -    scan all buffers for an overlap
    +    binary search for the buffers that can overlap
    ```
    
    - **`release_ranges` (commit 2):** finds the buffers that can overlap a
    released range with a binary search, as `PushBuffers::find` does. The
    trim and split logic does not change.
    - **`clear_ranges` (commit 3):** finds the buffers to remove with a
    binary search. Before, it built new buffer lists on each call. This
    commit gives most of the gain.
    - **`VecDeque` and `max_len_count` (commit 4):** `PushBuffers` stores
    the buffers in a `VecDeque`, so a removal at the start or at the end
    moves no other buffer. A count of the buffers with length `max_len` lets
    a removal skip the scan that updates `max_len` while other buffers of
    that length remain.
    
    <details>
    <summary>Benchmark numbers for each commit, and edge cases</summary>
    
    `push_decoder/prefetch` on the GKE runner. Each column adds one commit:
    
    | Buffers | Benchmark only | + `release_ranges` | + `clear_ranges` | +
    `VecDeque` |
    |--|--|--|--|--|
    | 1,000 | 1.12 ms | 1.09 ms | 0.64 ms | 0.60 ms |
    | 10,000 | 56.6 ms | 53.1 ms | 11.9 ms | 6.3 ms |
    | 30,000 | 478 ms | 447 ms | 73.4 ms | 19.2 ms |
    
    Change of the existing `1buf` and `Nbuf` cases for each commit:
    
    | Step | Change |
    |--|--|
    | Control (same code on both sides) | −3% to +2% |
    | + `release_ranges` | −1% to +2% |
    | + `clear_ranges` | −2% to +2% |
    | + `VecDeque` | +1% to +6% (slower in all seven cases) |
    
    Runs:
    
    - Benchmark only compared with `release_ranges`:
    https://github.com/apache/arrow-rs/pull/11459#issuecomment-6093342208
    - `release_ranges` compared with `clear_ranges`:
    https://github.com/apache/arrow-rs/pull/11459#issuecomment-6093337006
    - `clear_ranges` compared with `VecDeque`:
    https://github.com/apache/arrow-rs/pull/11459#issuecomment-6093332202
    - Control:
    https://github.com/apache/arrow-rs/pull/11459#issuecomment-6093343009
    
    These runs used the commits before the rebase onto `main`. The rebase
    changed only the benchmark file.
    
    Time for each `release_ranges` call, from temporary timing tests that
    are not part of this PR (release build, Apple silicon):
    
    | Buffers | Call | Before | After commit 2 |
    |--|--|--|--|
    | 1,000 | releases nothing | 495 ns | 55 ns |
    | 10,000 | releases nothing | 3.5 µs | 67 ns |
    | 100,000 | releases nothing | 37 µs | 61 ns |
    | 1,000 | trims one buffer | 1.5 µs | 82 ns |
    | 10,000 | trims one buffer | 8.5 µs | 91 ns |
    | 100,000 | trims one buffer | 86 µs | 81 ns |
    
    Edge cases:
    
    - Removal in the middle: `VecDeque::drain` moves the buffers on the
    shorter side of the removed buffers. Thus, only a removal at the start
    or at the end moves none.
    - `max_len`: it stays exact. It is computed again from all buffers only
    when the last buffer of length `max_len` is removed or trimmed.
    - Empty buffers: `push_range` accepts empty and inverted ranges, and
    `release_ranges` removes all of them when it releases bytes. To keep
    this behavior without a visit of all buffers, `PushBuffers` counts the
    empty buffers. `release_ranges` looks at all buffers only if the count
    is not zero.
    - Sort order: a trim moves the start of a buffer only past released
    bytes, so the trimmed buffers stay in order. Only split parts can be out
    of order, and the sort runs only after a split.
    - Buffers with the same start: they are adjacent but not sorted by end.
    `clear_ranges` checks each of them for an exact match.
    
    </details>
    
    # Are these changes tested?
    
    Yes. The existing tests cover this, in particular
    `fuzz_release_ranges_matches_rebuild`, which compares `release_ranges`
    with a full rebuild, and `fuzz_matches_a_linear_scan`, which includes
    random calls to `clear_ranges`. The new test
    `clear_ranges_with_buffers_of_the_same_start` covers buffers that share
    a start and ranges that are given more than once. `assert_invariants`
    now also checks the count of empty buffers and the count of buffers with
    length `max_len`.
    
    # Are there any user-facing changes?
    
    No. All changes are in `pub(crate)` functions and private fields, so a
    revert is safe if a regression shows.
    
    🤖 Generated with [Claude Code](https://claude.com/claude-code)
    
    ---------
    
    Co-authored-by: Claude Opus 5.5 <[email protected]>
---
 parquet/benches/push_decoder.rs  |  85 ++++++++++++-
 parquet/src/util/push_buffers.rs | 251 +++++++++++++++++++++++++++++----------
 2 files changed, 270 insertions(+), 66 deletions(-)

diff --git a/parquet/benches/push_decoder.rs b/parquet/benches/push_decoder.rs
index 72ad9f55dd..a66543f65b 100644
--- a/parquet/benches/push_decoder.rs
+++ b/parquet/benches/push_decoder.rs
@@ -24,6 +24,7 @@
 //! schemas with a page index.
 
 use std::hint::black_box;
+use std::ops::Range;
 use std::sync::Arc;
 
 use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
@@ -49,8 +50,16 @@ fn make_wide_schema(num_columns: usize) -> SchemaRef {
 
 /// Write a Parquet file with `num_columns` columns, 10 row groups of 100 rows.
 fn make_test_file(num_columns: usize) -> Bytes {
-    let num_rows = 1_000;
-    let rows_per_rg = 100;
+    make_test_file_with_row_groups(num_columns, 1_000, 100)
+}
+
+/// Write a Parquet file with `num_columns` columns and `num_rows` rows, in row
+/// groups of `rows_per_rg` rows.
+fn make_test_file_with_row_groups(
+    num_columns: usize,
+    num_rows: usize,
+    rows_per_rg: usize,
+) -> Bytes {
     let schema = make_wide_schema(num_columns);
     let columns: Vec<Arc<dyn arrow_array::Array>> = (0..num_columns)
         .map(|_| Arc::new(Float32Array::from(vec![0.0f32; num_rows])) as _)
@@ -129,6 +138,47 @@ fn build_readers_exact_ranges(
     }
 }
 
+/// Push one buffer per column chunk of the whole file before the first row
+/// group is read, then build all row group readers.
+///
+/// All buffers of the later row groups stay buffered while the decoder
+/// releases the bytes of each row group that it finished. Thus, this measures
+/// a release from many buffers.
+fn build_readers_prefetched(
+    file_data: &Bytes,
+    metadata: &Arc<parquet::file::metadata::ParquetMetaData>,
+) {
+    let mut decoder = 
ParquetPushDecoderBuilder::try_new_decoder(metadata.clone())
+        .unwrap()
+        .build()
+        .unwrap();
+
+    let ranges: Vec<Range<u64>> = metadata
+        .row_groups()
+        .iter()
+        .flat_map(|rg| rg.columns())
+        .map(|column| {
+            let (start, len) = column.byte_range();
+            start..start + len
+        })
+        .collect();
+    let buffers: Vec<Bytes> = ranges
+        .iter()
+        .map(|r| file_data.slice(r.start as usize..r.end as usize))
+        .collect();
+    decoder.push_ranges(ranges, buffers).unwrap();
+
+    loop {
+        match decoder.try_next_reader().unwrap() {
+            DecodeResult::Data(reader) => {
+                black_box(reader);
+            }
+            DecodeResult::Finished => break,
+            DecodeResult::NeedsData(r) => panic!("unexpected NeedsData: 
{r:?}"),
+        }
+    }
+}
+
 fn bench_1buf(c: &mut Criterion) {
     let mut group = c.benchmark_group("push_decoder/1buf");
 
@@ -173,6 +223,29 @@ fn bench_nbuf(c: &mut Criterion) {
     group.finish();
 }
 
+fn bench_prefetch(c: &mut Criterion) {
+    let mut group = c.benchmark_group("push_decoder/prefetch");
+
+    // 10 columns, one row in each row group.
+    for num_row_groups in [100, 1_000, 3_000] {
+        let file_data = make_test_file_with_row_groups(10, num_row_groups, 1);
+        let metadata = decode_metadata(&file_data);
+        let num_ranges: usize = metadata
+            .row_groups()
+            .iter()
+            .map(|rg| rg.columns().len())
+            .sum();
+
+        group.bench_with_input(
+            BenchmarkId::from_parameter(format!("{num_ranges}ranges")),
+            &(&file_data, &metadata),
+            |b, &(data, meta)| b.iter(|| build_readers_prefetched(data, meta)),
+        );
+    }
+
+    group.finish();
+}
+
 /// Write a Parquet file with `num_columns` columns and one row group of
 /// 1,000 rows, with `pages_per_column` data pages per column chunk.
 fn make_paged_test_file(num_columns: usize, pages_per_column: usize) -> Bytes {
@@ -251,5 +324,11 @@ fn bench_scan_plan(c: &mut Criterion) {
     group.finish();
 }
 
-criterion_group!(benches, bench_1buf, bench_nbuf, bench_scan_plan);
+criterion_group!(
+    benches,
+    bench_1buf,
+    bench_nbuf,
+    bench_prefetch,
+    bench_scan_plan
+);
 criterion_main!(benches);
diff --git a/parquet/src/util/push_buffers.rs b/parquet/src/util/push_buffers.rs
index a7c1675335..e1fa2bd049 100644
--- a/parquet/src/util/push_buffers.rs
+++ b/parquet/src/util/push_buffers.rs
@@ -18,6 +18,8 @@
 use crate::errors::ParquetError;
 use crate::file::reader::{ChunkReader, Length};
 use bytes::Bytes;
+use std::cmp::Ordering;
+use std::collections::VecDeque;
 use std::fmt::Display;
 use std::ops::Range;
 
@@ -54,10 +56,13 @@ pub struct PushBuffers {
     file_len: u64,
     /// The ranges of data that are available for decoding (not adjusted for
     /// offset), sorted by `start`
-    ranges: Vec<Range<u64>>,
+    ///
+    /// A `VecDeque`, because a scan pushes buffers at the end and releases
+    /// them at the start. Both are then `O(1)` for each buffer.
+    ranges: VecDeque<Range<u64>>,
     /// The buffers of data that can be used to decode the Parquet file, in the
     /// same order as `ranges`
-    buffers: Vec<Bytes>,
+    buffers: VecDeque<Bytes>,
     /// The length of the longest range in `ranges`.
     ///
     /// This keeps lookups fast in the common case: buffers that do not
@@ -66,9 +71,18 @@ 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 number of ranges in `ranges` with the length `max_len`. `max_len`
+    /// can decrease only when this count goes to zero. Thus, a release does
+    /// not scan all ranges to update `max_len` while other ranges of that
+    /// length remain.
+    max_len_count: usize,
     /// The sum of the lengths of `ranges`, kept up to date so that
     /// [`Self::buffered_bytes`] does not scan all ranges.
     buffered_bytes: u64,
+    /// The number of empty ranges in `ranges`. [`Self::push_range`] accepts
+    /// them, and [`Self::release_ranges`] removes them. The count tells
+    /// `release_ranges` when it must look for them in all buffers.
+    empty: usize,
 }
 
 impl Display for PushBuffers {
@@ -105,10 +119,12 @@ impl PushBuffers {
         Self {
             offset: 0,
             file_len,
-            ranges: Vec::new(),
-            buffers: Vec::new(),
+            ranges: VecDeque::new(),
+            buffers: VecDeque::new(),
             max_len: 0,
+            max_len_count: 0,
             buffered_bytes: 0,
+            empty: 0,
         }
     }
 
@@ -155,8 +171,13 @@ impl PushBuffers {
         // 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);
+        match expected.cmp(&self.max_len) {
+            Ordering::Greater => (self.max_len, self.max_len_count) = 
(expected, 1),
+            Ordering::Equal => self.max_len_count += 1,
+            Ordering::Less => {}
+        }
         self.buffered_bytes += expected;
+        self.empty += usize::from(expected == 0);
         self.ranges.insert(idx, range);
         self.buffers.insert(idx, buffer);
         Ok(())
@@ -197,8 +218,8 @@ impl PushBuffers {
         // long only if a caller pushes one large buffer and then many small
         // buffers after its start.
         let candidates = self.ranges.partition_point(|r| r.start <= start);
-        self.ranges[..candidates]
-            .iter()
+        self.ranges
+            .range(..candidates)
             .enumerate()
             .rev()
             .take_while(|(_, r)| r.start.saturating_add(self.max_len) >= end)
@@ -229,45 +250,77 @@ impl PushBuffers {
     }
 
     /// Clear any range and corresponding buffer that is exactly in the 
ranges_to_clear
+    ///
+    /// The cost is `O(log N)` in the number of buffers for each range to
+    /// clear.
     #[cfg(feature = "arrow")]
     pub(crate) fn clear_ranges(&mut self, ranges_to_clear: &[Range<u64>]) {
-        // Use `(start, end)` tuples because `Range` is not `Ord`
-        let mut clear: Vec<(u64, u64)> = ranges_to_clear.iter().map(|r| 
(r.start, r.end)).collect();
-        if clear.is_empty() {
-            return;
+        // The indexes of the buffers to remove. The buffers are sorted by
+        // start, so all buffers with the same start are adjacent.
+        let mut remove = vec![];
+        for clear in ranges_to_clear {
+            let first = self.ranges.partition_point(|r| r.start < clear.start);
+            let same_start = self
+                .ranges
+                .range(first..)
+                .take_while(|r| r.start == clear.start);
+            remove.extend(
+                same_start
+                    .enumerate()
+                    .filter(|(_, r)| r.end == clear.end)
+                    .map(|(idx, _)| first + idx),
+            );
         }
-        clear.sort_unstable();
-        let mut ranges = Vec::with_capacity(self.ranges.len());
-        let mut buffers = Vec::with_capacity(self.buffers.len());
-        for (range, buffer) in 
self.ranges.drain(..).zip(self.buffers.drain(..)) {
-            if clear.binary_search(&(range.start, range.end)).is_err() {
-                ranges.push(range);
-                buffers.push(buffer);
+        // Buffers with the same start are not sorted by end, and
+        // `ranges_to_clear` can contain a range more than once.
+        remove.sort_unstable();
+        remove.dedup();
+        let (Some(&first), Some(&last)) = (remove.first(), remove.last()) else 
{
+            return;
+        };
+        // Move the kept buffers in `first..=last` to the start of that span,
+        // in order, then remove the rest of the span. The buffers outside the
+        // span are not visited.
+        let mut kept = first;
+        let mut remove = remove.into_iter().peekable();
+        for idx in first..=last {
+            if remove.next_if_eq(&idx).is_some() {
+                let len = 
self.ranges[idx].end.saturating_sub(self.ranges[idx].start);
+                self.buffered_bytes -= len;
+                self.empty -= usize::from(len == 0);
+                if len == self.max_len {
+                    self.max_len_count -= 1;
+                }
+            } else {
+                self.ranges.swap(kept, idx);
+                self.buffers.swap(kept, idx);
+                kept += 1;
             }
         }
-        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();
+        self.ranges.drain(kept..=last);
+        self.buffers.drain(kept..=last);
+        // `max_len` changes only when no buffer of that length remains.
+        if self.max_len_count == 0 {
+            self.update_max_len();
+        }
     }
 
-    /// Set `max_len` to the maximum length of the remaining ranges.
+    /// Set `max_len` to the maximum length of the remaining ranges, and
+    /// `max_len_count` to the number of ranges of that length.
     ///
     /// A `max_len` that is too large is still correct, lookups only scan
     /// further. This update is for performance: a large buffer that was
     /// removed must not slow down later lookups.
     #[cfg(feature = "arrow")]
     fn update_max_len(&mut self) {
-        self.max_len = self
-            .ranges
-            .iter()
-            .map(|r| r.end.saturating_sub(r.start))
-            .max()
-            .unwrap_or(0);
+        (self.max_len, self.max_len_count) = (0, 0);
+        for len in self.ranges.iter().map(|r| r.end.saturating_sub(r.start)) {
+            match len.cmp(&self.max_len) {
+                Ordering::Greater => (self.max_len, self.max_len_count) = 
(len, 1),
+                Ordering::Equal => self.max_len_count += 1,
+                Ordering::Less => {}
+            }
+        }
     }
 
     /// Remove all buffered bytes in `ranges`.
@@ -281,38 +334,63 @@ impl PushBuffers {
     /// The buffers stay sorted by start, which lookups require (see
     /// [`Self::find`]). If no buffer overlaps `ranges`, the buffers do not
     /// change.
+    ///
+    /// The cost is `O(log N)` in the number of buffers for each range in
+    /// `ranges`, plus the number of buffers that overlap `ranges`.
     #[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;
+        // Find the buffers that overlap a released range. The buffers are
+        // sorted by start, and no buffer is longer than `max_len`. Thus, a
+        // buffer can overlap `r` only if it starts before `r.end` and less
+        // than `max_len` bytes before `r.start`:
+        //
+        //   buffers:  0..25    ├─────────┤
+        //             25..50             ├─────────┤
+        //             50..75                       ├─────────┤
+        //             75..100                                ├─────────┤
+        //   release:  55..70                         ├─────┤  max_len is 25: 
check 50..75 only
+        let mut overlapping = vec![];
+        let mut visited = 0;
+        for r in &release {
+            let max_len = self.max_len;
+            let first = self
+                .ranges
+                .partition_point(|b| b.start.saturating_add(max_len) <= 
r.start);
+            let end = self.ranges.partition_point(|b| b.start < r.end);
+            // A buffer that an earlier released range visited is not visited
+            // again, because `overlaps` checks all released ranges.
+            overlapping
+                .extend((first.max(visited)..end).filter(|&idx| 
overlaps(&self.ranges[idx])));
+            visited = visited.max(end);
         }
+        // Most calls release ranges that are no longer buffered. Return
+        // before changing the buffers.
+        let (Some(&first), Some(&last)) = (overlapping.first(), 
overlapping.last()) else {
+            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;
-            }
+        for idx in overlapping {
+            let (range, buffer) = (&mut self.ranges[idx], &mut 
self.buffers[idx]);
             let whole = range.clone();
+            // Each kept part is shorter than `whole`.
+            if whole.end - whole.start == self.max_len {
+                self.max_len_count -= 1;
+            }
             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);
@@ -351,30 +429,47 @@ impl PushBuffers {
                 }
             }
         }
-        if emptied {
+        if emptied || self.empty > 0 {
             // 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);
+            // buffers that were pushed. The released buffers are all in
+            // `first..=last`, so look at all buffers only if empty buffers
+            // were pushed.
+            let span = match self.empty {
+                0 => first..last + 1,
+                _ => 0..self.ranges.len(),
+            };
+            let mut kept = span.start;
+            for idx in span.clone() {
+                if !self.ranges[idx].is_empty() {
+                    self.ranges.swap(kept, idx);
+                    self.buffers.swap(kept, idx);
+                    kept += 1;
+                }
+            }
+            self.ranges.drain(kept..span.end);
+            self.buffers.drain(kept..span.end);
+            self.empty = 0;
         }
-        // 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) {
+        // A trim moves the start of a buffer only past released bytes. Each
+        // other buffer that starts in those bytes is trimmed past them too,
+        // or is empty and removed. Thus, the buffers are still sorted, except
+        // for the split parts, which are appended. Splits are rare, so sort
+        // all buffers then.
+        if !split.is_empty() {
             let mut parts: Vec<_> = std::mem::take(&mut self.ranges)
                 .into_iter()
                 .zip(std::mem::take(&mut self.buffers))
+                .chain(split)
                 .collect();
             parts.sort_by_key(|(range, _)| range.start);
             (self.ranges, self.buffers) = parts.into_iter().unzip();
         }
-        self.update_max_len();
+        // `max_len` changes only when no buffer of that length remains. Only
+        // then pay for the pass over all buffers, which keeps `max_len` exact
+        // (see `update_max_len`).
+        if self.max_len_count == 0 {
+            self.update_max_len();
+        }
     }
 
     /// Remove all buffered bytes outside `keep`.
@@ -398,7 +493,9 @@ impl PushBuffers {
         self.ranges.clear();
         self.buffers.clear();
         self.max_len = 0;
+        self.max_len_count = 0;
         self.buffered_bytes = 0;
+        self.empty = 0;
     }
 
     /// Panics if `ranges`, `buffers` and `max_len` do not agree, or if the
@@ -408,7 +505,7 @@ impl PushBuffers {
     fn assert_invariants(&self) {
         assert_eq!(self.ranges.len(), self.buffers.len());
         assert!(
-            self.ranges.is_sorted_by_key(|r| r.start),
+            self.ranges.iter().is_sorted_by_key(|r| r.start),
             "not sorted: {:?}",
             self.ranges
         );
@@ -420,7 +517,15 @@ impl PushBuffers {
             .iter()
             .map(|r| r.end.saturating_sub(r.start))
             .max();
-        assert_eq!(self.max_len, max_len.unwrap_or(0));
+        let max_len = max_len.unwrap_or(0);
+        assert_eq!(self.max_len, max_len);
+        let is_longest = |r: &&Range<u64>| r.end.saturating_sub(r.start) == 
max_len;
+        assert_eq!(
+            self.max_len_count,
+            self.ranges.iter().filter(is_longest).count()
+        );
+        let empty = self.ranges.iter().filter(|r| r.is_empty()).count();
+        assert_eq!(self.empty, empty);
     }
 }
 
@@ -712,6 +817,26 @@ mod tests {
         assert_eq!(buffers.ranges, vec![0..30, 10..20]);
     }
 
+    #[test]
+    #[cfg(feature = "arrow")]
+    fn clear_ranges_with_buffers_of_the_same_start() {
+        let mut buffers = PushBuffers::new(1000);
+        for range in [0..5, 10..20, 10..15, 10..20, 10..30, 10..12, 40..50] {
+            push(&mut buffers, range);
+        }
+        // Buffers with the same start are not sorted by end. A range to clear
+        // can match more than one buffer, and can be given more than once.
+        buffers.clear_ranges(&[10..30, 10..20, 10..20, 60..70]);
+        assert_valid(&buffers);
+        assert_eq!(buffers.ranges, vec![0..5, 10..15, 10..12, 40..50]);
+        assert_eq!(buffers.buffered_bytes(), 22);
+
+        // No buffer matches: the buffers do not change.
+        buffers.clear_ranges(&[10..20, 0..50]);
+        assert_valid(&buffers);
+        assert_eq!(buffers.ranges, vec![0..5, 10..15, 10..12, 40..50]);
+    }
+
     /// Random pushes, clears and lookups, compared with a list that is
     /// scanned in full for each lookup.
     #[test]
@@ -764,7 +889,7 @@ mod tests {
                     }
                 }
                 assert_valid(&buffers);
-                let mut actual = buffers.ranges.clone();
+                let mut actual = Vec::from(buffers.ranges.clone());
                 actual.sort_by_key(|r| (r.start, r.end));
                 model.sort_by_key(|r| (r.start, r.end));
                 assert_eq!(actual, model, "seed {seed}");

Reply via email to