From f4339952ce2544d4c92cd3e29e4bc1d82ab5cd86 Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Mon, 20 Jul 2026 11:37:01 +0000 Subject: [PATCH 01/12] feat(common): take_n_offsets helper A refinement of split_vec_min_alloc for splitting offset buffers. To improve performance of the current use in `BytesGroupValues::take_n`, this function adjusts the remaining offsets when they are copied to the new location, rather than two passes of copying then subtracting in place. It also accounts for n + 1 length of the resulting prefix vector, which current split_vec_min_alloc falls on the wrong side of. --- datafusion/common/src/utils/mod.rs | 47 ++++++++++++++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/datafusion/common/src/utils/mod.rs b/datafusion/common/src/utils/mod.rs index 94bbb91a7fa8b..4b39f8c8d80f9 100644 --- a/datafusion/common/src/utils/mod.rs +++ b/datafusion/common/src/utils/mod.rs @@ -50,6 +50,7 @@ use std::collections::HashSet; use std::iter::repeat_n; use std::num::NonZero; use std::ops::Range; +use std::ptr::copy_nonoverlapping; use std::sync::{Arc, LazyLock}; use std::thread::available_parallelism; @@ -411,6 +412,52 @@ pub fn split_vec_min_alloc(vec: &mut Vec, n: usize) -> Vec { } } +/// 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, vec.len() - n)`. This matters when the split emits a prefix +/// under memory pressure, where `n` can be close to `vec.len()`. +pub fn take_n_offsets(offsets: &mut Vec, n: usize) -> Vec +where + O: OffsetSizeTrait, +{ + let cut_offset = offsets[n]; + if n.saturating_mul(2) < offsets.len() { + let mut prefix = Vec::::with_capacity(n + 1); + // SAFETY: copying to a newly allocated vector with sufficient capacity. + // The length of `offsets` is checked to exceed n. + unsafe { + copy_nonoverlapping(offsets.as_ptr(), prefix.as_mut_ptr(), n); + *prefix.as_mut_ptr().add(n) = cut_offset; + prefix.set_len(n + 1); + } + // Shift the remaining offsets in place so that the first offset is 0. + let dst = offsets.as_mut_ptr(); + for (i, &offset) in offsets[n..].iter().enumerate() { + // SAFETY: the range of iteration is within `offsets.len()`. + // In the overlapping region, the destination element is overwritten + // only after it is read from. + unsafe { + *dst.add(i) = offset - cut_offset; + } + } + offsets.truncate(offsets.len() - n); + prefix + } else { + let remaining = offsets[n..] + .iter() + .map(|&offset| offset - cut_offset) + .collect::>(); + offsets.truncate(n + 1); + std::mem::replace(offsets, remaining) + } +} + #[cfg(test)] mod split_vec_min_alloc_tests { use super::split_vec_min_alloc; From 272964779c3765e8aaab7b550b5ece3e023d5dad Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Sat, 18 Jul 2026 11:41:44 +0000 Subject: [PATCH 02/12] bench: GroupValues::emit(First(n)) with bytes Benchmark the performance of take_n for ByteGroupValueBuilder in the context of the multi-group-by operation where it is used in production. --- .../physical-plan/benches/multi_group_by.rs | 194 +++++++++++++++++- 1 file changed, 193 insertions(+), 1 deletion(-) diff --git a/datafusion/physical-plan/benches/multi_group_by.rs b/datafusion/physical-plan/benches/multi_group_by.rs index 11c2800864316..2a2112374324a 100644 --- a/datafusion/physical-plan/benches/multi_group_by.rs +++ b/datafusion/physical-plan/benches/multi_group_by.rs @@ -27,11 +27,15 @@ //! covers a `(FixedSizeBinary, Int32)` key to exercise the //! `FixedSizeBinaryGroupValueBuilder`. -use arrow::array::{ArrayRef, Int32Array, UInt32Array}; +use arrow::array::{ + ArrayRef, BinaryArray, Int32Array, LargeBinaryArray, LargeStringArray, StringArray, + UInt32Array, +}; use arrow::compute::take; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use arrow::util::bench_util::create_fsb_array; use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use datafusion_expr::EmitTo; use datafusion_physical_plan::aggregates::group_values::GroupValues; use datafusion_physical_plan::aggregates::group_values::GroupValuesRows; use datafusion_physical_plan::aggregates::group_values::multi_group_by::GroupValuesColumn; @@ -444,6 +448,192 @@ fn bench_fixed_size_binary(c: &mut Criterion) { group.finish(); } +/// Offset width of the byte group columns used by the emit experiment. +/// +/// Both variants use two byte columns, so both route through +/// `ByteGroupValueBuilder` and its `take_n` (the code under test). `Large` uses +/// 64-bit offsets, which doubles the size of the offset vector that +/// `take_n_offsets` splits/shifts — making that path a larger share of `take_n`. +#[derive(Clone, Copy)] +enum OffsetWidth { + /// `(Binary, Utf8)` — 32-bit offsets. + Normal, + /// `(LargeBinary, LargeUtf8)` — 64-bit offsets. + Large, +} + +impl OffsetWidth { + fn label(self) -> &'static str { + match self { + OffsetWidth::Normal => "i32_offsets", + OffsetWidth::Large => "i64_offsets", + } + } + + /// The `(binary, utf8)` arrow types for this width. + fn data_types(self) -> (DataType, DataType) { + match self { + OffsetWidth::Normal => (DataType::Binary, DataType::Utf8), + OffsetWidth::Large => (DataType::LargeBinary, DataType::LargeUtf8), + } + } +} + +/// Schema for the emit experiment: a byte (`Binary`/`LargeBinary`) group column +/// paired with a string (`Utf8`/`LargeUtf8`) column. Both columns route the +/// multi-column GROUP BY through the `ByteGroupValueBuilder`, whose `take_n` +/// implementation is what the emit benchmark below exercises. +fn make_bytes_schema(width: OffsetWidth) -> SchemaRef { + let (binary, utf8) = width.data_types(); + Arc::new(Schema::new(vec![ + Field::new("b", binary, false), + Field::new("s", utf8, false), + ])) +} + +/// Generate byte-column batches with exactly `num_distinct_groups` distinct +/// keys, for the given [`OffsetWidth`]. Rows cycle through the group pool so the +/// accumulated group count is deterministic; values are short, variable-length +/// labels (`g0`, `g1`, …) — realistic keys without letting the value-buffer copy +/// dominate the offset work in `take_n`. The two columns carry identical labels +/// (as bytes / as str), so the combined cardinality equals `num_distinct_groups`. +fn generate_bytes_batches( + width: OffsetWidth, + num_distinct_groups: usize, + num_rows: usize, + batch_size: usize, +) -> Vec> { + let labels: Vec = (0..num_distinct_groups).map(|g| format!("g{g}")).collect(); + + let num_full_batches = num_rows / batch_size; + let remainder = num_rows % batch_size; + let num_batches = num_full_batches + if remainder > 0 { 1 } else { 0 }; + + (0..num_batches) + .map(|batch_idx| { + let batch_start = batch_idx * batch_size; + let current_batch_size = if batch_idx == num_batches - 1 && remainder > 0 { + remainder + } else { + batch_size + }; + + let group_ids = (0..current_batch_size) + .map(|row| (batch_start + row) % num_distinct_groups); + let strings = group_ids.clone().map(|g| labels[g].as_str()); + let bytes = group_ids.map(|g| g.to_be_bytes()); + + let (binary, utf8): (ArrayRef, ArrayRef) = match width { + OffsetWidth::Normal => ( + Arc::new(BinaryArray::from_iter_values(bytes)), + Arc::new(StringArray::from_iter_values(strings)), + ), + OffsetWidth::Large => ( + Arc::new(LargeBinaryArray::from_iter_values(bytes)), + Arc::new(LargeStringArray::from_iter_values(strings)), + ), + }; + vec![binary, utf8] + }) + .collect() +} + +/// Populate a fresh vectorized `GroupValuesColumn` by interning every batch, +/// returning it ready to emit. Used as the (untimed) setup for the emit bench. +/// +/// The builder's final state depends only on the set of distinct keys, not on +/// how many rows carry them — so to populate `g` groups it is enough to intern a +/// single batch of `g` distinct rows (see [`generate_bytes_batches`] with +/// `num_rows == num_distinct_groups`), which is what the small-cardinality +/// caller passes. +fn populate_group_values( + schema: &SchemaRef, + batches: &[Vec], +) -> Box { + let mut gv = create_group_values(schema, /* vectorized */ true); + let mut groups = Vec::with_capacity(DEFAULT_BATCH_SIZE); + for batch in batches { + groups.clear(); + gv.intern(batch, &mut groups).unwrap(); + } + gv +} + +/// The two `emit(EmitTo::First(n))` split-branch cases, keyed on `n` vs the +/// accumulated group count `g`. `take_n_offsets` +/// allocates for whichever side is smaller: +/// +/// * `shift_larger`: emit a smaller prefix — `n = g / 4`, shift remaining in place. +/// * `shift_mid`: emit the half-sized prefix in existing buffer — `n = g / 2`, +/// allocate and shift the remaining half. +/// * `shift_smaller`: emit a larger prefix in existing buffer — `n = g - g / 4`, +/// allocate and shift remaining. +fn emit_first_cases(num_groups: usize) -> [(&'static str, usize); 3] { + [ + ("shift_larger", num_groups / 4), + ("shift_midsize", num_groups / 2), + ("shift_smaller", num_groups - num_groups / 4), + ] +} + +/// Experiment 8a: `emit(EmitTo::First(n))` on a low-cardinality (1K groups) +/// `(Binary, Utf8)` key — reaches `ByteGroupValueBuilder::take_n`. +fn bench_emit_first_small(c: &mut Criterion) { + let mut group = c.benchmark_group("emit_first_small"); + group.sample_size(200); + + let num_groups = 1_000usize; + for width in [OffsetWidth::Normal, OffsetWidth::Large] { + let schema = make_bytes_schema(width); + let batches = + generate_bytes_batches(width, num_groups, num_groups, DEFAULT_BATCH_SIZE); + + for (branch, n) in emit_first_cases(num_groups) { + group.bench_with_input( + BenchmarkId::new(branch, format!("{}_grp_{num_groups}", width.label())), + &batches, + |b, batches| { + b.iter_batched_ref( + || populate_group_values(&schema, batches), + |gv| gv.emit(EmitTo::First(n)).unwrap(), + criterion::BatchSize::SmallInput, + ); + }, + ); + } + } + group.finish(); +} + +/// Experiment 8b: `emit(EmitTo::First(n))` on a high-cardinality (1M groups) +/// `(Binary, Utf8)` key — reaches `ByteGroupValueBuilder::take_n`. +fn bench_emit_first_large(c: &mut Criterion) { + let mut group = c.benchmark_group("emit_first_large"); + group.sample_size(50); + + let num_groups = 1_000_000usize; + for width in [OffsetWidth::Normal, OffsetWidth::Large] { + let schema = make_bytes_schema(width); + let batches = + generate_bytes_batches(width, num_groups, 1_000_000, DEFAULT_BATCH_SIZE); + + for (branch, n) in emit_first_cases(num_groups) { + group.bench_with_input( + BenchmarkId::new(branch, format!("{}_grp_{num_groups}", width.label())), + &batches, + |b, batches| { + b.iter_batched_ref( + || populate_group_values(&schema, batches), + |gv| gv.emit(EmitTo::First(n)).unwrap(), + criterion::BatchSize::LargeInput, + ); + }, + ); + } + } + group.finish(); +} + criterion_group!( benches, bench_issue_17850_regression, @@ -453,5 +643,7 @@ criterion_group!( bench_high_cardinality_scaling, bench_group_count_sweep, bench_fixed_size_binary, + bench_emit_first_small, + bench_emit_first_large, ); criterion_main!(benches); From 56a9d5a57ad26eae7af6f4f78a4b0624799a3e1e Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Tue, 21 Jul 2026 14:13:32 +0000 Subject: [PATCH 03/12] bench: drop emitted groups in emit_first Drop the returned groups from emit(First(n)) in the benchmarked function, to eliminate accumulating allocation behavior that's not characteristic of the use sites. --- datafusion/physical-plan/benches/multi_group_by.rs | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/datafusion/physical-plan/benches/multi_group_by.rs b/datafusion/physical-plan/benches/multi_group_by.rs index 2a2112374324a..ea32c25a066dc 100644 --- a/datafusion/physical-plan/benches/multi_group_by.rs +++ b/datafusion/physical-plan/benches/multi_group_by.rs @@ -595,7 +595,9 @@ fn bench_emit_first_small(c: &mut Criterion) { |b, batches| { b.iter_batched_ref( || populate_group_values(&schema, batches), - |gv| gv.emit(EmitTo::First(n)).unwrap(), + |gv| { + drop(black_box(gv.emit(EmitTo::First(n)).unwrap())); + }, criterion::BatchSize::SmallInput, ); }, @@ -624,7 +626,9 @@ fn bench_emit_first_large(c: &mut Criterion) { |b, batches| { b.iter_batched_ref( || populate_group_values(&schema, batches), - |gv| gv.emit(EmitTo::First(n)).unwrap(), + |gv| { + drop(black_box(gv.emit(EmitTo::First(n)).unwrap())); + }, criterion::BatchSize::LargeInput, ); }, From a1d95a2919cf7ce1a702bfde38861595c695955b Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Mon, 20 Jul 2026 11:45:48 +0000 Subject: [PATCH 04/12] perf(physical-plan): use take_n_offsets utility This is faster and could be reused with other GroupColumn implementations. --- .../aggregates/group_values/multi_group_by/bytes.rs | 11 ++--------- 1 file changed, 2 insertions(+), 9 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs index c83b1da4049bc..3e5a55c83792d 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs @@ -26,7 +26,7 @@ use arrow::array::{ use arrow::buffer::{OffsetBuffer, ScalarBuffer}; use arrow::datatypes::{ByteArrayType, DataType, GenericBinaryType}; use datafusion_common::utils::proxy::VecAllocExt; -use datafusion_common::utils::split_vec_min_alloc; +use datafusion_common::utils::take_n_offsets; use datafusion_common::{Result, exec_datafusion_err}; use datafusion_physical_expr_common::binary_map::{INITIAL_BUFFER_CAPACITY, OutputType}; use std::mem::size_of; @@ -378,14 +378,7 @@ where let null_buffer = self.nulls.take_n(n); let first_remaining_offset = O::as_usize(self.offsets[n]); - // Given offsets like [0, 2, 4, 5] and n = 1, we expect to get - // offsets [0, 2, 3]. We first create two offsets for first_n as [0, 2] and the remaining as [2, 4, 5]. - // And we shift the offset starting from 0 for the remaining one, [2, 4, 5] -> [0, 2, 3]. - let offset_n = self.offsets[n]; - let mut first_n_offsets = split_vec_min_alloc(&mut self.offsets, n); - // After the split, self.offsets[0] == offset_n in both branches; normalize in-place. - self.offsets.iter_mut().for_each(|o| *o = o.sub(offset_n)); - first_n_offsets.push(offset_n); + let first_n_offsets = take_n_offsets(&mut self.offsets, n); // SAFETY: the offsets were constructed correctly in `insert_if_new` -- // monotonically increasing, overflows were checked. From 3e6eb875d39920faf7acd8055224601a5694d82f Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Tue, 21 Jul 2026 15:05:39 +0000 Subject: [PATCH 05/12] docs: fix up argument name --- datafusion/common/src/utils/mod.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/datafusion/common/src/utils/mod.rs b/datafusion/common/src/utils/mod.rs index 4b39f8c8d80f9..ab3e90f82c201 100644 --- a/datafusion/common/src/utils/mod.rs +++ b/datafusion/common/src/utils/mod.rs @@ -420,8 +420,8 @@ pub fn split_vec_min_alloc(vec: &mut Vec, n: usize) -> Vec { /// in debug builds. /// /// Allocates for whichever side is smaller, so the new allocation is -/// `min(n + 1, vec.len() - n)`. This matters when the split emits a prefix -/// under memory pressure, where `n` can be close to `vec.len()`. +/// `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(offsets: &mut Vec, n: usize) -> Vec where O: OffsetSizeTrait, From 16319b87bb9758c12eb4fd4b25c288b1c14cdba8 Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Tue, 21 Jul 2026 23:14:15 +0300 Subject: [PATCH 06/12] refactor: safe way to copy out a Vec MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The loop does n + 1 iterations instead of n, but who's counting. Co-authored-by: Trương Hoàng Long --- datafusion/common/src/utils/mod.rs | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/datafusion/common/src/utils/mod.rs b/datafusion/common/src/utils/mod.rs index ab3e90f82c201..d124d6a49cc3f 100644 --- a/datafusion/common/src/utils/mod.rs +++ b/datafusion/common/src/utils/mod.rs @@ -428,14 +428,7 @@ where { let cut_offset = offsets[n]; if n.saturating_mul(2) < offsets.len() { - let mut prefix = Vec::::with_capacity(n + 1); - // SAFETY: copying to a newly allocated vector with sufficient capacity. - // The length of `offsets` is checked to exceed n. - unsafe { - copy_nonoverlapping(offsets.as_ptr(), prefix.as_mut_ptr(), n); - *prefix.as_mut_ptr().add(n) = cut_offset; - prefix.set_len(n + 1); - } + let prefix: Vec = offsets[..=n].into(); // Shift the remaining offsets in place so that the first offset is 0. let dst = offsets.as_mut_ptr(); for (i, &offset) in offsets[n..].iter().enumerate() { From d552beb7df87377a008b51a7474fe8e9fdd860b7 Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Tue, 21 Jul 2026 20:55:01 +0000 Subject: [PATCH 07/12] chore: remove unused import --- datafusion/common/src/utils/mod.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/datafusion/common/src/utils/mod.rs b/datafusion/common/src/utils/mod.rs index d124d6a49cc3f..d7642f4e6df5e 100644 --- a/datafusion/common/src/utils/mod.rs +++ b/datafusion/common/src/utils/mod.rs @@ -50,7 +50,6 @@ use std::collections::HashSet; use std::iter::repeat_n; use std::num::NonZero; use std::ops::Range; -use std::ptr::copy_nonoverlapping; use std::sync::{Arc, LazyLock}; use std::thread::available_parallelism; From 0803c90ba0c8f3fbe43512d0e62eae19663d9c12 Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Tue, 21 Jul 2026 20:55:27 +0000 Subject: [PATCH 08/12] refactor: no need for saturating_mul --- datafusion/common/src/utils/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/common/src/utils/mod.rs b/datafusion/common/src/utils/mod.rs index d7642f4e6df5e..f55f629f546b7 100644 --- a/datafusion/common/src/utils/mod.rs +++ b/datafusion/common/src/utils/mod.rs @@ -426,7 +426,7 @@ where O: OffsetSizeTrait, { let cut_offset = offsets[n]; - if n.saturating_mul(2) < offsets.len() { + if n < offsets.len() - n { let prefix: Vec = offsets[..=n].into(); // Shift the remaining offsets in place so that the first offset is 0. let dst = offsets.as_mut_ptr(); From c97e5d14ed9e3964ba5ae2e7c435c1ae8b414de5 Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Tue, 21 Jul 2026 22:06:57 +0000 Subject: [PATCH 09/12] bench(common): split_vec benchmark Created to test the new function take_n_offsets, but benchmarks for the similar split_vec_min_alloc can be added later. --- datafusion/common/Cargo.toml | 4 + datafusion/common/benches/split_vec.rs | 130 +++++++++++++++++++++++++ 2 files changed, 134 insertions(+) create mode 100644 datafusion/common/benches/split_vec.rs diff --git a/datafusion/common/Cargo.toml b/datafusion/common/Cargo.toml index 1eb23089a4021..f11a9865453ab 100644 --- a/datafusion/common/Cargo.toml +++ b/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 } diff --git a/datafusion/common/benches/split_vec.rs b/datafusion/common/benches/split_vec.rs new file mode 100644 index 0000000000000..8c54526cc3339 --- /dev/null +++ b/datafusion/common/benches/split_vec.rs @@ -0,0 +1,130 @@ +// 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` — 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 three 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`): +/// +/// * `shift_larger` (`n = g/4`): copy the smaller prefix out, shift the larger +/// remainder in place. +/// * `shift_midsize` (`n = g/2`): the balanced point; also the in-place-shift +/// branch, but with source/dest halves disjoint. +/// * `shift_smaller` (`n = g - g/4`): the other branch — allocate the smaller +/// remainder, return the larger prefix via `mem::replace`. +fn cases(num_groups: usize) -> [(&'static str, usize); 3] { + [ + ("shift_larger", num_groups / 4), + ("shift_midsize", num_groups / 2), + ("shift_smaller", 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(num_groups: usize) -> Vec { + (0..=num_groups).map(O::usize_as).collect() +} + +fn bench_take_n_offsets( + c: &mut Criterion, + group_name: &str, + num_groups: usize, + batch_size_hint: criterion::BatchSize, +) { + let mut group = c.benchmark_group(group_name); + + for width in [OffsetWidth::I32, OffsetWidth::I64] { + for (branch, n) in cases(num_groups) { + let id = + BenchmarkId::new(branch, format!("{}_grp_{num_groups}", 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::(num_groups), + |offsets| drop(black_box(take_n_offsets(offsets, n))), + batch_size_hint, + ), + OffsetWidth::I64 => b.iter_batched_ref( + || make_offsets::(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, + 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, + criterion::BatchSize::LargeInput, + ); +} + +criterion_group!( + benches, + bench_take_n_offsets_small, + bench_take_n_offsets_large +); +criterion_main!(benches); From 440e3a9e7bc00be695c71d550f8fb863ef91a1bb Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Tue, 21 Jul 2026 23:37:09 +0000 Subject: [PATCH 10/12] bench(common): tune take_n_offsets sample sizes --- datafusion/common/benches/split_vec.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/datafusion/common/benches/split_vec.rs b/datafusion/common/benches/split_vec.rs index 8c54526cc3339..b69ba06efd67e 100644 --- a/datafusion/common/benches/split_vec.rs +++ b/datafusion/common/benches/split_vec.rs @@ -71,9 +71,11 @@ 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 (branch, n) in cases(num_groups) { @@ -107,6 +109,7 @@ fn bench_take_n_offsets_small(c: &mut Criterion) { c, "take_n_offsets_small", 1_000, + 200, criterion::BatchSize::SmallInput, ); } @@ -118,6 +121,7 @@ fn bench_take_n_offsets_large(c: &mut Criterion) { c, "take_n_offsets_large", 1_000_000, + 50, criterion::BatchSize::LargeInput, ); } From f80450a3f5b85f71a4b63c420ba967a12fa023a7 Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Wed, 22 Jul 2026 09:57:33 +0000 Subject: [PATCH 11/12] bench: rationalize naming, add n == g / 2 + 1 case --- datafusion/common/benches/split_vec.rs | 25 +++++++++-------- .../physical-plan/benches/multi_group_by.rs | 28 ++++++++++--------- 2 files changed, 29 insertions(+), 24 deletions(-) diff --git a/datafusion/common/benches/split_vec.rs b/datafusion/common/benches/split_vec.rs index b69ba06efd67e..379e65513c9c4 100644 --- a/datafusion/common/benches/split_vec.rs +++ b/datafusion/common/benches/split_vec.rs @@ -43,21 +43,24 @@ impl OffsetWidth { } } -/// The three split-branch cases, keyed on `n` vs the number of groups `g`. +/// 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`): /// -/// * `shift_larger` (`n = g/4`): copy the smaller prefix out, shift the larger +/// * `n = g/4`: copy the smaller prefix out, shift the larger /// remainder in place. -/// * `shift_midsize` (`n = g/2`): the balanced point; also the in-place-shift -/// branch, but with source/dest halves disjoint. -/// * `shift_smaller` (`n = g - g/4`): the other branch — allocate the smaller +/// * `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) -> [(&'static str, usize); 3] { +fn cases(num_groups: usize) -> [usize; 4] { [ - ("shift_larger", num_groups / 4), - ("shift_midsize", num_groups / 2), - ("shift_smaller", num_groups - num_groups / 4), + num_groups / 4, + num_groups / 2, + num_groups / 2 + 1, + num_groups - num_groups / 4, ] } @@ -78,9 +81,9 @@ fn bench_take_n_offsets( group.sample_size(sample_size); for width in [OffsetWidth::I32, OffsetWidth::I64] { - for (branch, n) in cases(num_groups) { + for n in cases(num_groups) { let id = - BenchmarkId::new(branch, format!("{}_grp_{num_groups}", width.label())); + 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 diff --git a/datafusion/physical-plan/benches/multi_group_by.rs b/datafusion/physical-plan/benches/multi_group_by.rs index ea32c25a066dc..27283b56bdf56 100644 --- a/datafusion/physical-plan/benches/multi_group_by.rs +++ b/datafusion/physical-plan/benches/multi_group_by.rs @@ -559,20 +559,22 @@ fn populate_group_values( gv } -/// The two `emit(EmitTo::First(n))` split-branch cases, keyed on `n` vs the +/// The `emit(EmitTo::First(n))` split-branch cases, keyed on `n` vs the /// accumulated group count `g`. `take_n_offsets` /// allocates for whichever side is smaller: /// -/// * `shift_larger`: emit a smaller prefix — `n = g / 4`, shift remaining in place. -/// * `shift_mid`: emit the half-sized prefix in existing buffer — `n = g / 2`, -/// allocate and shift the remaining half. -/// * `shift_smaller`: emit a larger prefix in existing buffer — `n = g - g / 4`, +/// * `n = g / 4`: emit the smaller prefix, shift remaining in place. +/// * `n = g / 2`: emit the half-sized prefix, shift remaining in place. +/// * `n = g / 2 + 1`: emit the just-over-half-sized prefix in existing buffer, /// allocate and shift remaining. -fn emit_first_cases(num_groups: usize) -> [(&'static str, usize); 3] { +/// * `n = g - g / 4`: emit the larger prefix in existing buffer, +/// allocate and shift remaining. +fn emit_first_cases(num_groups: usize) -> [usize; 4] { [ - ("shift_larger", num_groups / 4), - ("shift_midsize", num_groups / 2), - ("shift_smaller", num_groups - num_groups / 4), + num_groups / 4, + num_groups / 2, + num_groups / 2 + 1, + num_groups - num_groups / 4, ] } @@ -588,9 +590,9 @@ fn bench_emit_first_small(c: &mut Criterion) { let batches = generate_bytes_batches(width, num_groups, num_groups, DEFAULT_BATCH_SIZE); - for (branch, n) in emit_first_cases(num_groups) { + for n in emit_first_cases(num_groups) { group.bench_with_input( - BenchmarkId::new(branch, format!("{}_grp_{num_groups}", width.label())), + BenchmarkId::new(format!("grp_{num_groups}_emit_{n}",), width.label()), &batches, |b, batches| { b.iter_batched_ref( @@ -619,9 +621,9 @@ fn bench_emit_first_large(c: &mut Criterion) { let batches = generate_bytes_batches(width, num_groups, 1_000_000, DEFAULT_BATCH_SIZE); - for (branch, n) in emit_first_cases(num_groups) { + for n in emit_first_cases(num_groups) { group.bench_with_input( - BenchmarkId::new(branch, format!("{}_grp_{num_groups}", width.label())), + BenchmarkId::new(format!("grp_{num_groups}_emit_{n}",), width.label()), &batches, |b, batches| { b.iter_batched_ref( From 2b0cc933bbd5d58dfad0d1e4d72f12275e09cded Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Wed, 22 Jul 2026 23:46:19 +0300 Subject: [PATCH 12/12] refactor: rewrite unsafe loop In the take_n_offsets drain-and-shift branch. The rewritten, safe version vectorizes just as well and performs the same in the benchmarks. --- datafusion/common/src/utils/mod.rs | 100 +++++++++++++++++++++++++---- 1 file changed, 89 insertions(+), 11 deletions(-) diff --git a/datafusion/common/src/utils/mod.rs b/datafusion/common/src/utils/mod.rs index f55f629f546b7..daf92e6fd3639 100644 --- a/datafusion/common/src/utils/mod.rs +++ b/datafusion/common/src/utils/mod.rs @@ -427,18 +427,14 @@ where { let cut_offset = offsets[n]; if n < offsets.len() - n { - let prefix: Vec = offsets[..=n].into(); - // Shift the remaining offsets in place so that the first offset is 0. - let dst = offsets.as_mut_ptr(); - for (i, &offset) in offsets[n..].iter().enumerate() { - // SAFETY: the range of iteration is within `offsets.len()`. - // In the overlapping region, the destination element is overwritten - // only after it is read from. - unsafe { - *dst.add(i) = offset - cut_offset; - } + 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(offsets.len() - n); + offsets.truncate(new_len); prefix } else { let remaining = offsets[n..] @@ -565,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, Vec) { + 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 = (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 ///