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-11742-01a9917658fecd639991e92a5a5810be710b539c in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
commit c4ef83d86a2ad0531bda67e9b61d9477c6a85bd8 Author: Neil Conway <[email protected]> AuthorDate: Thu Oct 22 15:43:18 2026 +0000 bench: add `MutableBuffer::extend`, enhance `MutableArrayData` (#11742) # Which issue does this PR close? - N/A. Split out of #11656, which uses these benchmarks. # Rationale for this change Extend the benchmarks for `MutableBuffer::extend` to cover more iterator shapes, and extend the `MutableArrayData` benchmarks to cover more data types. We also revise how data is generated for the `zip_kernels` and `mutable_array` benchmarks: the previous approach repeatedly passed the same input data to the functions on every iteration. This allows the branch predictor to do much better than it would on real-world data, which initially hid the ARM regression described above. These improvements are intended to evaluate #11656 more accurately. # What changes are included in this PR? - Add an arrow-buffer `mutable_buffer_extend` benchmark that appends 32Ki items, one `extend` call per chunk, with chunks of random length (16 or 1024 items on average). It covers copied and mapped slices, an `i32` range, `repeat_n`, a `Chain` of two slices, and a `Filter`, whose size hint is not exact. - Add `Utf8View` and dictionary cases to `mutable_array`, with the existing long random slices and with slices of 1 to 8 rows. Every case now starts with room for all the rows it copies, so its time goes to the copies rather than to growing the buffers. - Cycle `zip_kernels`, `mutable_array` and `mutable_buffer_extend` through 16 to 64 inputs instead of reusing one. Scalars and the constant masks stay the same, as they would across the batches of a query. # Are these changes tested? Benchmark-only change. # Are there any significant user-facing changes? No. # AI usage Claude Code (Opus 5.5), reviewed with Codex (Astra 6.1). Co-authored-by: Ed Seidl <[email protected]> --- arrow-buffer/Cargo.toml | 4 + arrow-buffer/benches/mutable_buffer_extend.rs | 136 ++++++++++++++++++++++++++ arrow/benches/mutable_array.rs | 108 +++++++++++++++++--- arrow/benches/zip_kernels.rs | 80 ++++++++++----- 4 files changed, 289 insertions(+), 39 deletions(-) diff --git a/arrow-buffer/Cargo.toml b/arrow-buffer/Cargo.toml index 465f4eee48..39bdf82120 100644 --- a/arrow-buffer/Cargo.toml +++ b/arrow-buffer/Cargo.toml @@ -64,5 +64,9 @@ harness = false name = "mutable_buffer_repeat_slice" harness = false +[[bench]] +name = "mutable_buffer_extend" +harness = false + [lints] workspace = true diff --git a/arrow-buffer/benches/mutable_buffer_extend.rs b/arrow-buffer/benches/mutable_buffer_extend.rs new file mode 100644 index 0000000000..5ba77f95f3 --- /dev/null +++ b/arrow-buffer/benches/mutable_buffer_extend.rs @@ -0,0 +1,136 @@ +// 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 appending iterators to a [`MutableBuffer`] with `extend`. +//! +//! Each case appends `LEN` items to a buffer that already has enough +//! capacity, with one `extend` call per chunk of items. The chunks have random +//! lengths that average `chunk` items. + +use arrow_buffer::{ArrowNativeType, MutableBuffer}; +use criterion::*; +use rand::rngs::StdRng; +use rand::{RngExt, SeedableRng}; +use std::hint; +use std::ops::Range; + +const LEN: usize = 1 << 15; + +/// Each case cycles through this many inputs, one per iteration. An input is a +/// split of `0..LEN` into chunks, plus an array of items for cases that read +/// them. +const NUM_INPUTS: usize = 16; + +/// `NUM_INPUTS` arrays of `LEN` items made by `item` +fn random_inputs<T>(mut item: impl FnMut() -> T) -> Vec<Vec<T>> { + (0..NUM_INPUTS) + .map(|_| (0..LEN).map(|_| item()).collect()) + .collect() +} + +/// For each input, `0..LEN` split into chunks whose lengths are chosen at random +/// from 1 to `2 * chunk - 1` items +fn random_chunks(rng: &mut StdRng, chunk: usize) -> Vec<Vec<Range<usize>>> { + (0..NUM_INPUTS) + .map(|_| { + let mut chunks = vec![]; + let mut start = 0; + while start < LEN { + let end = LEN.min(start + rng.random_range(1..2 * chunk)); + chunks.push(start..end); + start = end; + } + chunks + }) + .collect() +} + +/// Benchmarks `extend` with the iterators that `make` returns. Each iteration +/// takes the next input, clears the buffer, and appends `make(input, range)` +/// for each `range` in `chunks[input]`. +fn bench_extend<I: Iterator<Item: ArrowNativeType>>( + group: &mut BenchmarkGroup<'_, measurement::WallTime>, + name: &str, + chunk: usize, + chunks: &[Vec<Range<usize>>], + make: impl Fn(usize, Range<usize>) -> I, +) { + let mut buffer = MutableBuffer::new(LEN * size_of::<I::Item>()); + let mut inputs = (0..NUM_INPUTS).cycle(); + group.bench_function(BenchmarkId::new(name, chunk), |b| { + b.iter(|| { + let input = inputs.next().unwrap(); + buffer.clear(); + for range in &chunks[input] { + buffer.extend(make(input, range.clone())); + } + hint::black_box(buffer.as_slice()); + }) + }); +} + +fn criterion_benchmark(c: &mut Criterion) { + let mut rng = StdRng::seed_from_u64(42); + let bytes = random_inputs(|| rng.random_range(b'a'..=b'z')); + let keys: Vec<Vec<i32>> = random_inputs(|| rng.random_range(0..1000)); + let wide: Vec<Vec<u128>> = random_inputs(|| rng.random_range(0..1 << 100)); + + // Read at run time so the compiler can't specialize the maps for them. + let (from, to) = hint::black_box((b'a', b'z')); + let key_offset: i32 = hint::black_box(1000); + let wide_offset: u128 = hint::black_box(1000); + + let mut group = c.benchmark_group("MutableBuffer extend"); + for chunk in [16, 1024] { + let chunks = random_chunks(&mut rng, chunk); + bench_extend(&mut group, "copied u8", chunk, &chunks, |input, r| { + bytes[input][r].iter().copied() + }); + // Like replacing one byte with another in a string. + bench_extend(&mut group, "map u8", chunk, &chunks, |input, r| { + bytes[input][r] + .iter() + .map(move |&b| if b == from { to } else { b }) + }); + bench_extend(&mut group, "map i32", chunk, &chunks, |input, r| { + keys[input][r].iter().map(move |k| k + key_offset) + }); + bench_extend(&mut group, "map u128", chunk, &chunks, |input, r| { + wide[input][r].iter().map(move |v| v + wide_offset) + }); + bench_extend(&mut group, "range i32", chunk, &chunks, |_, r| { + r.start as i32..r.end as i32 + }); + bench_extend(&mut group, "repeat_n u32", chunk, &chunks, |_, r| { + std::iter::repeat_n(r.start as u32, r.len()) + }); + bench_extend(&mut group, "chain i32", chunk, &chunks, |input, r| { + let mid = r.len() / 2; + let (lo, hi) = keys[input][r].split_at(mid); + lo.iter().chain(hi).copied() + }); + // The size hint is not exact. The predicate always passes, so the + // branch is predictable. + bench_extend(&mut group, "filter i32", chunk, &chunks, |input, r| { + keys[input][r].iter().copied().filter(|&k| k >= 0) + }); + } + group.finish(); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/arrow/benches/mutable_array.rs b/arrow/benches/mutable_array.rs index e06cc446dd..f1b9fce6c9 100644 --- a/arrow/benches/mutable_array.rs +++ b/arrow/benches/mutable_array.rs @@ -19,13 +19,34 @@ extern crate criterion; use criterion::Criterion; -use rand::RngExt; +use std::sync::Arc; -use arrow::util::test_util::seedable_rng; -use arrow::{array::*, util::bench_util::create_string_array}; +use rand::rngs::StdRng; +use rand::{RngExt, SeedableRng}; -fn create_slices(size: usize) -> Vec<(usize, usize)> { - let rng = &mut seedable_rng(); +use arrow::array::*; +use arrow::datatypes::Int32Type; +use arrow::util::bench_util::{ + create_string_array_with_len_range_and_prefix_and_seed, + create_string_view_array_with_len_range_and_seed, +}; + +/// Each benchmark cycles through this many arrays and lists of slices, one per +/// iteration. +const NUM_INPUTS: usize = 64; + +/// `NUM_INPUTS` values made by `make` from the seeds `0..NUM_INPUTS` +fn per_seed<T>(make: impl FnMut(u64) -> T) -> Vec<T> { + (0..NUM_INPUTS as u64).map(make).collect() +} + +/// `size` random strings of 0 to 399 bytes +fn create_strings(size: usize, null_density: f32, seed: u64) -> StringArray { + create_string_array_with_len_range_and_prefix_and_seed(size, null_density, 0, 399, "", seed) +} + +fn create_slices(size: usize, seed: u64) -> Vec<(usize, usize)> { + let rng = &mut StdRng::seed_from_u64(seed); (0..size) .map(|_| { @@ -36,23 +57,84 @@ fn create_slices(size: usize) -> Vec<(usize, usize)> { .collect() } -fn bench<T: Array>(v1: &T, slices: &[(usize, usize)]) { +/// Slices of 1 to 8 rows, like the runs that `zip` copies for a `CASE` +/// expression whose condition changes often. +fn create_short_slices(size: usize, seed: u64) -> Vec<(usize, usize)> { + let rng = &mut StdRng::seed_from_u64(seed); + + (0..size) + .map(|_| { + let start = rng.random_range(0..size - 8); + let len = rng.random_range(1..=8); + (start, start + len) + }) + .collect() +} + +/// A dictionary array of `size` random keys into 16 random strings +fn create_dict_array(size: usize, seed: u64) -> DictionaryArray<Int32Type> { + let rng = &mut StdRng::seed_from_u64(seed); + let keys = Int32Array::from_iter_values((0..size).map(|_| rng.random_range(0..16))); + DictionaryArray::new(keys, Arc::new(create_strings(16, 0.0, seed))) +} + +/// Total number of rows in `slices` +fn total_len(slices: &[(usize, usize)]) -> usize { + slices.iter().map(|(start, end)| end - start).sum() +} + +/// Copies `slices` of `v1`, which hold `len` rows in all, into a +/// `MutableArrayData` that starts with room for all of them. +fn bench<T: Array>(v1: &T, slices: &[(usize, usize)], len: usize) { let data = v1.to_data(); - let mut mutable = MutableArrayData::new(vec![&data], false, 5); + let mut mutable = MutableArrayData::new(vec![&data], false, len); for (start, end) in slices { mutable.try_extend(0, *start, *end).unwrap(); } mutable.freeze(); } +/// Benchmarks [`bench`], where iteration `i` copies list `i % NUM_INPUTS` of +/// `slices` from array `i % NUM_INPUTS` of `arrays`. +fn bench_inputs<T: Array>( + c: &mut Criterion, + name: &str, + arrays: &[T], + slices: &[Vec<(usize, usize)>], +) { + let lens: Vec<usize> = slices.iter().map(|slices| total_len(slices)).collect(); + let mut inputs = arrays.iter().zip(slices).zip(&lens).cycle(); + c.bench_function(name, |b| { + b.iter(|| { + let ((array, slices), len) = inputs.next().unwrap(); + bench(array, slices, *len) + }) + }); +} + fn add_benchmark(c: &mut Criterion) { - let v1 = create_string_array::<i32>(1024, 0.0); - let v2 = create_slices(1024); - c.bench_function("mutable str 1024", |b| b.iter(|| bench(&v1, &v2))); + let long_slices = per_seed(|seed| create_slices(1024, seed)); + let short_slices = per_seed(|seed| create_short_slices(1024, seed)); + + let strings = per_seed(|seed| create_strings(1024, 0.0, seed)); + bench_inputs(c, "mutable str 1024", &strings, &long_slices); + + let strings = per_seed(|seed| create_strings(1024, 0.5, seed)); + bench_inputs(c, "mutable str nulls 1024", &strings, &long_slices); + + let views = + per_seed(|seed| create_string_view_array_with_len_range_and_seed(1024, 0.0, 0..400, seed)); + bench_inputs(c, "mutable str_view 1024", &views, &long_slices); + bench_inputs( + c, + "mutable str_view short slices 1024", + &views, + &short_slices, + ); - let v1 = create_string_array::<i32>(1024, 0.5); - let v2 = create_slices(1024); - c.bench_function("mutable str nulls 1024", |b| b.iter(|| bench(&v1, &v2))); + let dicts = per_seed(|seed| create_dict_array(1024, seed)); + bench_inputs(c, "mutable dict 1024", &dicts, &long_slices); + bench_inputs(c, "mutable dict short slices 1024", &dicts, &short_slices); } criterion_group!(benches, add_benchmark); diff --git a/arrow/benches/zip_kernels.rs b/arrow/benches/zip_kernels.rs index 3d29baded3..dbf6884331 100644 --- a/arrow/benches/zip_kernels.rs +++ b/arrow/benches/zip_kernels.rs @@ -29,6 +29,10 @@ use arrow::datatypes::*; use arrow::util::bench_util::*; use arrow_select::zip::zip; +/// Each benchmark cycles through this many arrays and random masks, one per +/// iteration. +const NUM_INPUTS: usize = 64; + trait InputGenerator { fn name(&self) -> &str; @@ -163,20 +167,27 @@ impl InputGenerator for GenerateStringView { } } -fn mask_cases(len: usize) -> Vec<(&'static str, BooleanArray)> { +/// Returns the masks for each mask case: `NUM_INPUTS` of each random case, and +/// one of each constant case. +fn mask_cases(len: usize) -> Vec<(&'static str, Vec<BooleanArray>)> { + let random = |null_density, true_density| -> Vec<BooleanArray> { + (0..NUM_INPUTS as u64) + .map(|seed| create_boolean_array_with_seed(len, null_density, true_density, seed)) + .collect() + }; vec![ - ("all_true", create_boolean_array(len, 0.0, 1.0)), - ("99pct_true", create_boolean_array(len, 0.0, 0.99)), - ("90pct_true", create_boolean_array(len, 0.0, 0.9)), - ("50pct_true", create_boolean_array(len, 0.0, 0.5)), + ("all_true", vec![create_boolean_array(len, 0.0, 1.0)]), + ("99pct_true", random(0.0, 0.99)), + ("90pct_true", random(0.0, 0.9)), + ("50pct_true", random(0.0, 0.5)), ( "true_then_false", - BooleanArray::from_iter((0..len).map(|i| i < len / 2)), + vec![BooleanArray::from_iter((0..len).map(|i| i < len / 2))], ), - ("10pct_true", create_boolean_array(len, 0.0, 0.1)), - ("1pct_true", create_boolean_array(len, 0.0, 0.01)), - ("all_false", create_boolean_array(len, 0.0, 0.0)), - ("50pct_nulls", create_boolean_array(len, 0.5, 0.5)), + ("10pct_true", random(0.0, 0.1)), + ("1pct_true", random(0.0, 0.01)), + ("all_false", vec![create_boolean_array(len, 0.0, 0.0)]), + ("50pct_nulls", random(0.5, 0.5)), ] } @@ -192,8 +203,13 @@ fn bench_zip_on_input_generator(c: &mut Criterion, input_generator: &impl InputG .try_into() .unwrap(); - let array_1_10pct_nulls = input_generator.generate_array(42, ARRAY_LEN, 0.1); - let array_2_10pct_nulls = input_generator.generate_array(18, ARRAY_LEN, 0.1); + let arrays = |first_seed: u64| -> Vec<ArrayRef> { + (0..NUM_INPUTS as u64) + .map(|i| input_generator.generate_array(first_seed + i, ARRAY_LEN, 0.1)) + .collect() + }; + let arrays_1_10pct_nulls = arrays(42); + let arrays_2_10pct_nulls = arrays(42 + NUM_INPUTS as u64); let masks = mask_cases(ARRAY_LEN); @@ -206,8 +222,8 @@ fn bench_zip_on_input_generator(c: &mut Criterion, input_generator: &impl InputG description, &mut group, &masks, - &Scalar::new(truthy), - &Scalar::new(falsy), + &[Scalar::new(truthy)], + &[Scalar::new(falsy)], ); } @@ -215,40 +231,52 @@ fn bench_zip_on_input_generator(c: &mut Criterion, input_generator: &impl InputG "array_vs_non_null_scalar", &mut group, &masks, - &array_1_10pct_nulls, - &Scalar::new(non_null_scalar_1.clone()), + &arrays_1_10pct_nulls, + &[Scalar::new(non_null_scalar_1.clone())], ); bench_zip_input_on_all_masks( "non_null_scalar_vs_array", &mut group, &masks, - &Scalar::new(non_null_scalar_1.clone()), - &array_1_10pct_nulls, + &[Scalar::new(non_null_scalar_1.clone())], + &arrays_1_10pct_nulls, ); bench_zip_input_on_all_masks( "array_vs_array", &mut group, &masks, - &array_1_10pct_nulls, - &array_2_10pct_nulls, + &arrays_1_10pct_nulls, + &arrays_2_10pct_nulls, ); group.finish(); } +/// Benchmarks `zip` for each mask case. Each iteration takes the next of the +/// case's masks, of `truthy` and of `falsy`, cycling through each list on its +/// own, so a one-element list stays the same. fn bench_zip_input_on_all_masks( description: &str, group: &mut BenchmarkGroup<WallTime>, - masks: &[(&str, BooleanArray)], - truthy: &impl Datum, - falsy: &impl Datum, + mask_cases: &[(&str, Vec<BooleanArray>)], + truthy: &[impl Datum], + falsy: &[impl Datum], ) { - for (mask_description, mask) in masks { + for (mask_description, masks) in mask_cases { let id = BenchmarkId::new(description, mask_description); - group.bench_with_input(id, mask, |b, mask| { - b.iter(|| hint::black_box(zip(mask, truthy, falsy))) + let mut masks = masks.iter().cycle(); + let mut truthy = truthy.iter().cycle(); + let mut falsy = falsy.iter().cycle(); + group.bench_function(id, |b| { + b.iter(|| { + hint::black_box(zip( + masks.next().unwrap(), + truthy.next().unwrap(), + falsy.next().unwrap(), + )) + }) }); } }
