perf: improve splitting of the offset vector in GroupColumn::take_n implementations, adding a new helper fn

datafusion/common/Cargo.toml

@@ -64,6 +64,10 @@ name = "scalar_to_array"
 harness = false
 name = "stats_merge"
 
+[[bench]]
+harness = false
+name = "split_vec"
+
 [dependencies]
 arrow = { workspace = true }
 arrow-ipc = { workspace = true }

datafusion/common/benches/split_vec.rs

@@ -0,0 +1,137 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Benchmarks for [`take_n_offsets`], which splits an offset vector at index
+//! `n` (returning `[offsets[0..=n]]` and leaving the shifted remainder in
+//! place). It underpins `ByteGroupValueBuilder::take_n`, but is measured here in
+//! isolation on a plain `Vec<O>` — no arrays, no builder, no `emit` — so the
+//! offset-splitting work is the whole signal.
+
+use arrow::array::OffsetSizeTrait;
+use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
+use datafusion_common::utils::take_n_offsets;
+use std::hint::black_box;
+
+/// Offset width under test: 32-bit (`Utf8`/`Binary`) vs 64-bit
+/// (`LargeUtf8`/`LargeBinary`) offsets. 64-bit doubles the bytes moved.
+#[derive(Clone, Copy)]
+enum OffsetWidth {
+    I32,
+    I64,
+}
+
+impl OffsetWidth {
+    fn label(self) -> &'static str {
+        match self {
+            OffsetWidth::I32 => "i32_offsets",
+            OffsetWidth::I64 => "i64_offsets",
+        }
+    }
+}
+
+/// The split-branch cases, keyed on `n` vs the number of groups `g`.
+/// `take_n_offsets` allocates for whichever side is smaller (predicate
+/// `n < offsets.len() - n`, where `offsets.len() == g + 1`):
+///
+/// * `n = g/4`: copy the smaller prefix out, shift the larger
+///   remainder in place.
+/// * `n = g/2`: the half-point; also the in-place-shift
+///   branch.
+/// * `n = g/2 + 1`: one item beyond the half-point,
+///   to see if it works better on the other side of the fence.
+/// * `n = g - g/4`: the other branch — allocate the smaller
+///   remainder, return the larger prefix via `mem::replace`.
+fn cases(num_groups: usize) -> [usize; 4] {
+    [
+        num_groups / 4,
+        num_groups / 2,
+        num_groups / 2 + 1,
+        num_groups - num_groups / 4,
+    ]
+}
+
+/// A monotonic offset vector of length `num_groups + 1` (`0, 1, …, g`), matching
+/// the shape a byte builder holds for `g` groups.
+fn make_offsets<O: OffsetSizeTrait>(num_groups: usize) -> Vec<O> {
+    (0..=num_groups).map(O::usize_as).collect()
+}
+
+fn bench_take_n_offsets(
+    c: &mut Criterion,
+    group_name: &str,
+    num_groups: usize,
+    sample_size: usize,
+    batch_size_hint: criterion::BatchSize,
+) {
+    let mut group = c.benchmark_group(group_name);
+    group.sample_size(sample_size);
+
+    for width in [OffsetWidth::I32, OffsetWidth::I64] {
+        for n in cases(num_groups) {
+            let id =
+                BenchmarkId::new(format!("grp_{num_groups}_take_{n}"), width.label());
+            group.bench_function(id, |b| {
+                // `take_n_offsets` mutates its input, so rebuild a fresh offset
+                // vector per iteration in the untimed setup. Monomorphize per
+                // width; only the `take_n_offsets` call is timed.
+                match width {
+                    OffsetWidth::I32 => b.iter_batched_ref(
+                        || make_offsets::<i32>(num_groups),
+                        |offsets| drop(black_box(take_n_offsets(offsets, n))),
+                        batch_size_hint,
+                    ),
+                    OffsetWidth::I64 => b.iter_batched_ref(
+                        || make_offsets::<i64>(num_groups),
+                        |offsets| drop(black_box(take_n_offsets(offsets, n))),
+                        batch_size_hint,
+                    ),
+                }
+            });
+        }
+    }
+    group.finish();
+}
+
+/// Low-cardinality (1K groups): offset vectors fit in cache.
+fn bench_take_n_offsets_small(c: &mut Criterion) {
+    bench_take_n_offsets(
+        c,
+        "take_n_offsets_small",
+        1_000,
+        200,
+        criterion::BatchSize::SmallInput,
+    );
+}
+
+/// High-cardinality (1M groups): offset vectors exceed cache, so the split's
+/// memory-access pattern dominates.
+fn bench_take_n_offsets_large(c: &mut Criterion) {
+    bench_take_n_offsets(
+        c,
+        "take_n_offsets_large",
+        1_000_000,
+        50,
+        criterion::BatchSize::LargeInput,
+    );
+}
+
+criterion_group!(
+    benches,
+    bench_take_n_offsets_small,
+    bench_take_n_offsets_large
+);
+criterion_main!(benches);

datafusion/common/src/utils/mod.rs

@@ -411,6 +411,41 @@ pub fn split_vec_min_alloc<T>(vec: &mut Vec<T>, n: usize) -> Vec<T> {
     }
 }
 
+/// Splits a vector of offsets at index `n`, returning a vector with
+/// `[offsets[0], ..., offsets[n]]` and leaving the remaining offsets in shifted
+/// `offsets`, adjusted to start at 0. This function assumes monotonicity of the
+/// `offsets` elements, so the `offsets[n]` cut-off value is subtracted from
+/// the remaining offset values with no overflow checks, except those enabled
+/// in debug builds.
+///
+/// Allocates for whichever side is smaller, so the new allocation is
+/// `min(n + 1, offsets.len() - n)`. This matters when the split emits a prefix
+/// under memory pressure, where `n` can be close to `offsets.len()`.
+pub fn take_n_offsets<O>(offsets: &mut Vec<O>, n: usize) -> Vec<O>
+where
+    O: OffsetSizeTrait,
+{
+    let cut_offset = offsets[n];
+    if n < offsets.len() - n {
+        let prefix = offsets[..=n].to_vec();
+        let new_len = offsets.len() - n;
+        // Shift the retained tail down to the front and rebase it to start at 0
+        // in a single fused pass.
+        for i in 0..new_len {
+            offsets[i] = offsets[n + i] - cut_offset;
+        }
+        offsets.truncate(new_len);
+        prefix
+    } else {
+        let remaining = offsets[n..]
+            .iter()
+            .map(|&offset| offset - cut_offset)
+            .collect::<Vec<O>>();
+        offsets.truncate(n + 1);
+        std::mem::replace(offsets, remaining)
+    }
+}
+
 #[cfg(test)]
 mod split_vec_min_alloc_tests {
     use super::split_vec_min_alloc;
@@ -526,6 +561,88 @@ mod split_vec_min_alloc_tests {
     }
 }
 
+#[cfg(test)]
+mod take_n_offsets_tests {
+    use super::take_n_offsets;
+
+    /// Reference implementation used to cross-check both branches: the emitted
+    /// prefix is `offsets[..=n]`, and the retained offsets are `offsets[n..]`
+    /// rebased so the first is 0.
+    fn expected(offsets: &[i64], n: usize) -> (Vec<i64>, Vec<i64>) {
+        let cut = offsets[n];
+        let prefix = offsets[..=n].to_vec();
+        let remaining = offsets[n..].iter().map(|&o| o - cut).collect();
+        (prefix, remaining)
+    }
+
+    fn check(offsets: &[i64], n: usize) {
+        let (exp_prefix, exp_remaining) = expected(offsets, n);
+        let mut v = offsets.to_vec();
+        let prefix = take_n_offsets(&mut v, n);
+        assert_eq!(prefix, exp_prefix, "prefix mismatch for n={n}");
+        assert_eq!(v, exp_remaining, "remaining mismatch for n={n}");
+    }
+
+    #[test]
+    fn in_place_shift_branch() {
+        // n < len - n  ->  in-place shift branch (branch A). len = 11, n = 3:
+        // new_len = 8 > n = 3, so the write range [0..8) overlaps the read
+        // range [3..11) -- exercises the overlapping-copy correctness.
+        let offsets = [0, 2, 5, 9, 14, 20, 27, 35, 44, 54, 65];
+        check(&offsets, 3);
+    }
+
+    #[test]
+    fn allocate_remaining_branch() {
+        // n >= len - n  ->  allocate-remaining branch (branch B). len = 11, n = 8.
+        let offsets = [0, 2, 5, 9, 14, 20, 27, 35, 44, 54, 65];
+        check(&offsets, 8);
+    }
+
+    #[test]
+    fn pivot_boundary() {
+        // Straddle the `n < len - n` pivot the way the benchmarks do (g/2 and
+        // g/2 + 1). With len = 2m + 1, n = m takes branch A and n = m + 1 takes
+        // branch B.
+        let m = 500usize;
+        let offsets: Vec<i64> = (0..=(2 * m as i64)).map(|i| i * 3).collect();
+        assert_eq!(offsets.len(), 2 * m + 1);
+        check(&offsets, m); // branch A: n=500, new_len=501 > n
+        check(&offsets, m + 1); // branch B: n=501, new_len=500 <= n
+    }
+
+    #[test]
+    fn take_one() {
+        // Smallest branch-A case: n = 1, so the shift loop's src index n + i
+        // always leads the dst index i by exactly 1.
+        let offsets = [0, 4, 7, 11, 18];
+        check(&offsets, 1);
+    }
+
+    #[test]
+    fn take_all_but_last() {
+        // Largest branch-B case: n = len - 1, remaining is a single 0.
+        let offsets = [0, 4, 7, 11, 18];
+        check(&offsets, offsets.len() - 1);
+    }
+
+    #[test]
+    fn i32_offsets() {
+        // Confirm monomorphization over the other OffsetSizeTrait width, both
+        // branches.
+        let offsets: [i32; 7] = [0, 3, 6, 10, 15, 21, 28];
+        let mut a = offsets.to_vec();
+        let pa = take_n_offsets(&mut a, 2); // branch A
+        assert_eq!(pa, vec![0, 3, 6]);
+        assert_eq!(a, vec![0, 4, 9, 15, 22]);
+
+        let mut b = offsets.to_vec();
+        let pb = take_n_offsets(&mut b, 5); // branch B
+        assert_eq!(pb, vec![0, 3, 6, 10, 15, 21]);
+        assert_eq!(b, vec![0, 7]);
+    }
+}
+
 /// Creates single element [`ListArray`], [`LargeListArray`] and
 /// [`FixedSizeListArray`] from other arrays
 ///