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..379e65513c9c4 --- /dev/null +++ b/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` — 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(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, + 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::(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, + 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); diff --git a/datafusion/common/src/utils/mod.rs b/datafusion/common/src/utils/mod.rs index 94bbb91a7fa8b..daf92e6fd3639 100644 --- a/datafusion/common/src/utils/mod.rs +++ b/datafusion/common/src/utils/mod.rs @@ -411,6 +411,41 @@ 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, 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, +{ + 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::>(); + 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, 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 /// diff --git a/datafusion/physical-plan/benches/multi_group_by.rs b/datafusion/physical-plan/benches/multi_group_by.rs index 11c2800864316..27283b56bdf56 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,198 @@ 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 `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: +/// +/// * `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. +/// * `n = g - g / 4`: emit the larger prefix in existing buffer, +/// allocate and shift remaining. +fn emit_first_cases(num_groups: usize) -> [usize; 4] { + [ + num_groups / 4, + num_groups / 2, + num_groups / 2 + 1, + 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 n in emit_first_cases(num_groups) { + group.bench_with_input( + BenchmarkId::new(format!("grp_{num_groups}_emit_{n}",), width.label()), + &batches, + |b, batches| { + b.iter_batched_ref( + || populate_group_values(&schema, batches), + |gv| { + drop(black_box(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 n in emit_first_cases(num_groups) { + group.bench_with_input( + BenchmarkId::new(format!("grp_{num_groups}_emit_{n}",), width.label()), + &batches, + |b, batches| { + b.iter_batched_ref( + || populate_group_values(&schema, batches), + |gv| { + drop(black_box(gv.emit(EmitTo::First(n)).unwrap())); + }, + criterion::BatchSize::LargeInput, + ); + }, + ); + } + } + group.finish(); +} + criterion_group!( benches, bench_issue_17850_regression, @@ -453,5 +649,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); 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.