From 1f80e4c3e331b3622fdd695193d9263d92348399 Mon Sep 17 00:00:00 2001 From: Neil Conway Date: Thu, 15 Oct 2026 08:19:05 -0400 Subject: [PATCH] bench: add `MutableBuffer::extend` and view/dictionary `MutableArrayData` benchmarks, with new inputs each iteration `MutableBuffer::extend` had one benchmark, which copies `u32` slices. Add an arrow-buffer benchmark that appends 32Ki items, one `extend` call per chunk, with chunks of 16 or 1024 items on average. It covers these iterator shapes: - copied slices of `u8`, and mapped slices of `u8`, `i32` and `u128`. - an `i32` range. - `repeat_n`. - a `Chain` of two slices. - a `Filter`, whose size hint, unlike the others', is not exact. `mutable_array` only covered `Utf8`. Add `Utf8View` and dictionary cases, 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 from 5 rows. `zip_kernels` and `mutable_array` ran every iteration on the same input. The CPU's branch predictor then learns that input's exact sequence of branch outcomes, such as which string views are stored inline or how long each copied run is. A query sees each batch once, so it never gets that benefit, and code that branches on the data looks faster than it would be in a query. On an AMD EPYC Milan server, zipping two arrays of 8192 string views of 3 to 19 bytes with an all-true mask (one copy, with a branch per view on whether it is inline) took about 3 branch misses per call when every call used the same arrays. Cycling through 32 or more pairs of arrays raised that to about 3,400, and the call took 3.4x as long. The predictor memorized up to 32Ki to 64Ki views, depending on the build, and almost none at 128Ki. With lists of 1024 slices of 1 to 8 rows, it memorized up to 8 lists; from 32 lists on, it missed about once per slice. Every data input now changes each iteration, cycling through: - `zip_kernels`: 64 arrays for each array input and 64 masks for each random mask case (512Ki rows). - `mutable_array`: 64 arrays and 64 lists of slices. - `mutable_buffer_extend`: 16 splits of the items into chunks, and 16 arrays of 32Ki items for each case that reads input data (512Ki items). Scalars and the constant masks stay the same on every iteration, as they would across the batches of a query. `mutable_buffer_extend` chunks have random lengths, 1 to 31 or 1 to 2047 items. With fixed lengths, the predictor would learn when each call's loop ends, and a loop that handles several items per pass would never have items left over. --- 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(-) create mode 100644 arrow-buffer/benches/mutable_buffer_extend.rs diff --git a/arrow-buffer/Cargo.toml b/arrow-buffer/Cargo.toml index 465f4eee4898..39bdf82120b4 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 000000000000..5ba77f95f3cd --- /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(mut item: impl FnMut() -> T) -> Vec> { + (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>> { + (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>( + group: &mut BenchmarkGroup<'_, measurement::WallTime>, + name: &str, + chunk: usize, + chunks: &[Vec>], + make: impl Fn(usize, Range) -> I, +) { + let mut buffer = MutableBuffer::new(LEN * size_of::()); + 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> = random_inputs(|| rng.random_range(0..1000)); + let wide: Vec> = 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 e06cc446dde7..f1b9fce6c98f 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(make: impl FnMut(u64) -> T) -> Vec { + (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(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 { + 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(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( + c: &mut Criterion, + name: &str, + arrays: &[T], + slices: &[Vec<(usize, usize)>], +) { + let lens: Vec = 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::(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::(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 3d29baded370..dbf6884331eb 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)> { + let random = |null_density, true_density| -> Vec { + (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 { + (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, - masks: &[(&str, BooleanArray)], - truthy: &impl Datum, - falsy: &impl Datum, + mask_cases: &[(&str, Vec)], + 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(), + )) + }) }); } }