From cfab89fcc165122e48ce1ccbc4a0638eb57fee65 Mon Sep 17 00:00:00 2001 From: Neil Conway Date: Thu, 15 Oct 2026 08:19:05 -0400 Subject: [PATCH 1/5] 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(), + )) + }) }); } } From e2eecf12101913f94f1ec573951f4b5bdc645c47 Mon Sep 17 00:00:00 2001 From: Neil Conway Date: Fri, 16 Oct 2026 23:15:05 -0400 Subject: [PATCH 2/5] perf: faster `MutableBuffer::extend` for exact-size iterators MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `MutableBuffer::extend` checks the remaining capacity before writing each item, then pushes any items left once the capacity runs out. When the iterator's size hint gives an exact length, reserve that many items and zip the iterator with `0..len` instead. For slice iterators and other iterators that support random access, `zip` then walks both sides by index, and the compiler can often vectorize the loop. Iterators without an exact size hint, such as `Filter`, keep the existing loop. The new loop sets the buffer's length from the loop index after each item. Adding the item size to the length instead made the compiler keep a running total in vector registers: on x86 the vectorized loop then handled half as many `i32` items per pass, and 1024-item calls on mapped `i32` slices were 38% to 45% slower than before. `MutableArrayData` copies string views and dictionary keys through this method, and so do `BufferBuilder::append_n` and collecting into a `MutableBuffer`. Its string view copy returned inline views (12 bytes or fewer) unchanged and added the buffer offset to the others. On an Apple M4 Max, the new loop turned that choice into a branch per view, which the CPU often guesses wrong when short and long strings are mixed, and zipping string view arrays took up to 3.3x as long as before. The copy now adds 0 to inline views instead, so every view takes the same path. The existing loop already had that branch on x86, so view copies get faster there too. Benchmarks, each version run twice in alternating order. Changes are against the benchmark commit. The two runs of a version were within 2% of each other for over 90% of the cases; larger gaps that matter for a result are noted. - AMD EPYC Milan: - `zip_kernels` string views: -16% to -69%. With an all-true or all-false mask, a zip took 9 µs instead of 28 µs, and its branch misses per 8192 views fell from 3,464 to 13. The two 50% masks with arrays on both sides, which `zip` copies with `interleave`, are -2%. `Utf8` is within 1.5%. - `mutable_array`: `str_view` -18% with short slices and +12% with long slices, whose strings of 0 to 399 bytes are almost all longer than 12 bytes, so the CPU always guessed the branch right. `dict` -10% with long slices and -14% with short ones. `Utf8` unchanged: with nulls, one run without this change took 29 ms, while the other three, and four more runs of each version made separately, took 35.5 to 37 ms. - `mutable_buffer_extend`: calls of 16 items -9% to -31%, except `Filter` +2%. Calls of 1024 items: mapped slices -5% (`i32`, with runs 5% apart) to -21%, copied `u8` -6%, the others within 3%. - Apple M4 Max: - `zip_kernels` string views: -3% to -22%. The two 50% masks with arrays on both sides are within 1%. `Utf8` is within 2.1%. - `mutable_array`: `str_view` -17% with short slices and -6% with long slices. `dict` -14% with long slices and -31% with short ones. `Utf8` unchanged. - `mutable_buffer_extend`: calls of 16 items -22% to -57%, except copied `u8` (-5%, with runs 14% apart) and `Filter` (+11%, 0.4 to 1.4 ns per call; slower in every run with this change, including earlier runs of it). Calls of 1024 items: mapped `u8` -31%, mapped `u128` -19%, `Chain` -9% (with runs 6% apart), the others within 5% except `Filter`, whose runs were up to 24% apart. An iterator whose size hint reports an exact length but that yields more items now has the extra items dropped; before, they were appended. Such an iterator breaks the `Iterator::size_hint` contract. An iterator that yields fewer items than it reports still has all of them appended. A lower bound whose size in bytes overflows `usize` is now ignored. Before, the multiplication panicked in debug builds and wrapped in release builds. --- arrow-buffer/src/buffer/mutable.rs | 121 +++++++++++++++++++++++++++-- arrow-data/src/transform/mod.rs | 11 ++- 2 files changed, 120 insertions(+), 12 deletions(-) diff --git a/arrow-buffer/src/buffer/mutable.rs b/arrow-buffer/src/buffer/mutable.rs index fc641009ae34..9524885ccff2 100644 --- a/arrow-buffer/src/buffer/mutable.rs +++ b/arrow-buffer/src/buffer/mutable.rs @@ -995,11 +995,20 @@ impl MutableBuffer { mut iterator: I, ) { let item_size = std::mem::size_of::(); - let (lower, _) = iterator.size_hint(); - let additional = lower * item_size; - self.reserve(additional); + let (lower, upper) = iterator.size_hint(); + if let Some(byte_count) = lower.checked_mul(item_size) { + self.reserve(byte_count); + if upper == Some(lower) { + // SAFETY: `reserve` just made room for `lower` items + return unsafe { self.extend_from_exact_size_iter(iterator, lower) }; + } + } - // this is necessary because of https://github.com/rust-lang/rust/issues/32155 + // Keep the length in a local and write it back when `len` is dropped. + // If the loop updated `self.len` directly, the compiler couldn't rule + // out that the writes through `dst` change it, so it would load and + // store the length for every item. `Vec` works around the same problem + // (https://github.com/rust-lang/rust/issues/32155). let mut len = SetLenOnDrop::new(&mut self.len); let mut dst = unsafe { self.data.as_ptr().add(len.local_len) }; let capacity = self.layout.size(); @@ -1021,8 +1030,48 @@ impl MutableBuffer { iterator.for_each(|item| self.push(item)); } + /// Appends the items of `iterator`, whose size hint says it yields exactly + /// `count` items. + /// + /// The loop doesn't check the capacity before each item. If `iterator` + /// yields more than `count` items, the extra items are not appended. + /// + /// # Safety + /// + /// The buffer must have room for `count` more items. + #[inline] + unsafe fn extend_from_exact_size_iter>( + &mut self, + iterator: I, + count: usize, + ) { + debug_assert!( + count + .checked_mul(std::mem::size_of::()) + .is_some_and(|bytes| bytes <= self.capacity() - self.len), + "buffer has no room for {count} more items" + ); + + // Keep the length in a local and write it back when `len` is dropped. + // If the loop updated `self.len` directly, the compiler couldn't rule + // out that the writes through `dst` change it, so it would load and + // store the length for every item. `Vec` works around the same problem + // (https://github.com/rust-lang/rust/issues/32155). + let mut len = SetLenOnDrop::new(&mut self.len); + let start = len.local_len; + // SAFETY: `start` is at most the capacity. It need not be a multiple of + // `T`'s alignment, so `dst` is only written unaligned. + let dst = unsafe { self.data.as_ptr().add(start) }.cast::(); + for (i, item) in (0..count).zip(iterator) { + // SAFETY: `i < count`, and the caller guarantees room for `count` + // items after the buffer's length. + unsafe { dst.add(i).write_unaligned(item) }; + len.local_len = start + (i + 1) * std::mem::size_of::(); + } + } + /// Creates a [`MutableBuffer`] from an [`Iterator`] with a trusted (upper) length. - /// Prefer this to `collect` whenever possible, as it is faster ~60% faster. + /// Prefer this to `collect` whenever possible, as it is ~60% faster. /// # Example /// ``` /// # use arrow_buffer::buffer::MutableBuffer; @@ -1105,7 +1154,7 @@ impl MutableBuffer { /// Creates a [`MutableBuffer`] from an [`Iterator`] with a trusted (upper) length or errors /// if any of the items of the iterator is an error. - /// Prefer this to `collect` whenever possible, as it is faster ~60% faster. + /// Prefer this to `collect` whenever possible, as it is ~60% faster. /// /// # Errors /// @@ -1366,6 +1415,66 @@ mod tests { assert_eq!(&[1u8, 1u8, 0, 0, 0, 0, 0, 0, 0], buf.as_slice()); } + /// Yields `items`, but its size hint always claims exactly `claimed` items. + struct WrongSizeHint<'a> { + items: std::slice::Iter<'a, u32>, + claimed: usize, + } + + impl Iterator for WrongSizeHint<'_> { + type Item = u32; + + fn next(&mut self) -> Option { + self.items.next().copied() + } + + fn size_hint(&self) -> (usize, Option) { + (self.claimed, Some(self.claimed)) + } + } + + /// `extend` is safe, so an exact size hint that is wrong must not make it + /// write past the reserved space or claim bytes it never wrote. Miri + /// catches both when it runs this test. + #[test] + fn mutable_extend_from_iter_wrong_exact_size_hint() { + // Fewer items than claimed: the length covers only the items written + let mut buf = MutableBuffer::new(0); + buf.extend(WrongSizeHint { + items: [1, 2].iter(), + claimed: 3, + }); + assert_eq!(buf.typed_data::(), &[1, 2]); + + // More items than claimed: nothing is written past the reserved space + let items: Vec = (0..32).collect(); + let mut buf = MutableBuffer::new(0); + buf.extend(WrongSizeHint { + items: items.iter(), + claimed: 2, + }); + assert_eq!(buf.typed_data::(), &[0, 1]); + + // A claimed size whose byte count overflows `usize`: the hint is ignored + let mut buf = MutableBuffer::new(0); + buf.extend(WrongSizeHint { + items: [1, 2].iter(), + claimed: usize::MAX / 4 + 1, + }); + assert_eq!(buf.typed_data::(), &[1, 2]); + } + + #[test] + fn mutable_extend_from_iter_inexact_size_hint() { + // `filter` reports a lower bound of 0, so nothing is reserved up + // front: the first 16 items fill the existing capacity and the rest + // are pushed + let mut buf = MutableBuffer::new(64); + buf.extend((0..100_u32).filter(|x| x % 3 == 0)); + let expected: Vec = (0..100).filter(|x| x % 3 == 0).collect(); + assert_eq!(buf.typed_data::(), &expected[..]); + } + #[test] fn mutable_extend_from_slice_unaligned_u64() { let mut buf = MutableBuffer::new(16); diff --git a/arrow-data/src/transform/mod.rs b/arrow-data/src/transform/mod.rs index 2db03a42aa9b..a7255b0254bd 100644 --- a/arrow-data/src/transform/mod.rs +++ b/arrow-data/src/transform/mod.rs @@ -226,13 +226,12 @@ fn build_extend_view(array: &ArrayData, buffer_offset: u32) -> Extend<'_> { mutable .buffer1 .extend(views[start..start + len].iter().map(|v| { - let len = *v as u32; - if len <= 12 { - return *v; // Stored inline - } + // A view of 12 bytes or fewer stores the string inline and + // has no buffer index, so adding 0 leaves it unchanged. + let offset = if *v as u32 > 12 { buffer_offset } else { 0 }; let mut view = ByteView::from(*v); - view.buffer_index += buffer_offset; - view.into() + view.buffer_index += offset; + u128::from(view) })); Ok(()) }, From 9b988c5b696d464dd180329bd1f8643d9752b98f Mon Sep 17 00:00:00 2001 From: Neil Conway Date: Mon, 19 Oct 2026 12:00:14 -0400 Subject: [PATCH 3/5] perf: add the buffer offset to string views with a shift `MutableArrayData` adds a buffer offset to the buffer index of each string view it copies, and adds 0 to views of inline strings. It unpacked each view into a `ByteView`, added the offset to the `buffer_index` field, and packed the fields back into a `u128`. Now it shifts the offset to the position of the buffer index, bits 64 to 95, and adds it to the whole view. The new index is less than the number of data buffers, so the sum never carries into the bits above it. After the previous commit, x86 copies of long runs of views that all hold buffer indices were up to 18% slower than before it. A null-free column of strings longer than 12 bytes has only such views. The shifted add makes those copies 12% to 23% faster than before the previous commit. Benchmarks on an AMD EPYC Milan server, each version run twice in alternating order. Changes are against the benchmark commit, then against the previous commit: - `mutable_array` `str_view` with long slices, whose strings of 0 to 399 bytes are almost all longer than 12 bytes: -23% (the previous commit: +10%), -30% against the previous commit. With short slices: -12%, +3% against the previous commit. - `zip_kernels` string views, two arrays: -36% to -77%, and -27% to +1% against the previous commit. The exception is 10% true rows, where one of the two runs of this commit was 11% slower than the other; three more runs of each version put this commit 1% faster than the previous one. - `merge` of null-free string view arrays, with a benchmark that is not part of this change: every case is 10% to 81% faster. The seven cases that the previous commit made 7.5% to 18% slower, copies of 13- to 40-byte strings in long runs, are 12% to 16% faster. Against the previous commit: -27% to +2.6%. - `take` on lists of 0 to 8 string views, also with a benchmark that is not part of this change: -38%, +2.4% against the previous commit. --- arrow-data/src/transform/mod.rs | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/arrow-data/src/transform/mod.rs b/arrow-data/src/transform/mod.rs index a7255b0254bd..f5f6cf63990a 100644 --- a/arrow-data/src/transform/mod.rs +++ b/arrow-data/src/transform/mod.rs @@ -20,7 +20,7 @@ //! Provides utilities for creating, manipulating, and converting Arrow arrays //! made of primitive types, strings, and nested types. -use super::{ArrayData, ArrayDataBuilder, ByteView, data::new_buffers}; +use super::{ArrayData, ArrayDataBuilder, data::new_buffers}; use crate::bit_mask::set_bits; use arrow_buffer::buffer::{BooleanBuffer, NullBuffer}; use arrow_buffer::{ArrowNativeType, Buffer, IntervalMonthDayNano, MutableBuffer, bit_util, i256}; @@ -229,9 +229,10 @@ fn build_extend_view(array: &ArrayData, buffer_offset: u32) -> Extend<'_> { // A view of 12 bytes or fewer stores the string inline and // has no buffer index, so adding 0 leaves it unchanged. let offset = if *v as u32 > 12 { buffer_offset } else { 0 }; - let mut view = ByteView::from(*v); - view.buffer_index += offset; - u128::from(view) + // The buffer index is bits 64 to 95 of a view. The new index + // is less than the number of data buffers, so the sum never + // carries into the bits above it. + v + ((offset as u128) << 64) })); Ok(()) }, From 2edf0bbf3ca52dd7e0a8a56bea980f82a5817f3e Mon Sep 17 00:00:00 2001 From: Neil Conway Date: Mon, 19 Oct 2026 13:08:48 -0400 Subject: [PATCH 4/5] chore: address review feedback on `extend` and string view copies When an iterator's size hint reports an exact length, `MutableBuffer::extend` appends at most that many items, so it drops the extra items of an iterator whose size hint is wrong. The `Extend` and `FromIterator` impls of `MutableBuffer` and `BufferBuilder` all append this way, and their docs now say so. The string view copy compares a view's length with `MAX_INLINE_VIEW_LEN` instead of a literal 12, as `ByteView` does. --- arrow-buffer/src/buffer/mutable.rs | 8 ++++++++ arrow-buffer/src/builder/mod.rs | 8 ++++++++ arrow-data/src/transform/mod.rs | 8 ++++++-- 3 files changed, 22 insertions(+), 2 deletions(-) diff --git a/arrow-buffer/src/buffer/mutable.rs b/arrow-buffer/src/buffer/mutable.rs index 9524885ccff2..d506a3f5bff2 100644 --- a/arrow-buffer/src/buffer/mutable.rs +++ b/arrow-buffer/src/buffer/mutable.rs @@ -960,6 +960,10 @@ pub(crate) fn dangling_ptr() -> NonNull { } impl Extend for MutableBuffer { + /// Appends the items of `iter`. + /// + /// If `iter`'s [`size_hint`](Iterator::size_hint) reports an exact length, + /// at most that many items are appended. #[inline] fn extend>(&mut self, iter: T) { let iterator = iter.into_iter(); @@ -1341,6 +1345,10 @@ impl std::iter::FromIterator for MutableBuffer { } impl std::iter::FromIterator for MutableBuffer { + /// Collects the items of `iter`. + /// + /// If `iter`'s [`size_hint`](Iterator::size_hint) reports an exact length, + /// at most that many items are collected. fn from_iter>(iter: I) -> Self { let mut buffer = Self::default(); buffer.extend_from_iter(iter.into_iter()); diff --git a/arrow-buffer/src/builder/mod.rs b/arrow-buffer/src/builder/mod.rs index 289374359a5a..11959a5ffe03 100644 --- a/arrow-buffer/src/builder/mod.rs +++ b/arrow-buffer/src/builder/mod.rs @@ -376,6 +376,10 @@ impl Default for BufferBuilder { } impl Extend for BufferBuilder { + /// Appends the items of `iter`. + /// + /// If `iter`'s [`size_hint`](Iterator::size_hint) reports an exact length, + /// at most that many items are appended. fn extend>(&mut self, iter: I) { self.buffer.extend(iter) } @@ -391,6 +395,10 @@ impl From> for BufferBuilder { } impl FromIterator for BufferBuilder { + /// Collects the items of `iter`. + /// + /// If `iter`'s [`size_hint`](Iterator::size_hint) reports an exact length, + /// at most that many items are collected. fn from_iter>(iter: I) -> Self { let mut builder = Self::default(); builder.extend(iter); diff --git a/arrow-data/src/transform/mod.rs b/arrow-data/src/transform/mod.rs index f5f6cf63990a..f72d781e9546 100644 --- a/arrow-data/src/transform/mod.rs +++ b/arrow-data/src/transform/mod.rs @@ -20,7 +20,7 @@ //! Provides utilities for creating, manipulating, and converting Arrow arrays //! made of primitive types, strings, and nested types. -use super::{ArrayData, ArrayDataBuilder, data::new_buffers}; +use super::{ArrayData, ArrayDataBuilder, MAX_INLINE_VIEW_LEN, data::new_buffers}; use crate::bit_mask::set_bits; use arrow_buffer::buffer::{BooleanBuffer, NullBuffer}; use arrow_buffer::{ArrowNativeType, Buffer, IntervalMonthDayNano, MutableBuffer, bit_util, i256}; @@ -228,7 +228,11 @@ fn build_extend_view(array: &ArrayData, buffer_offset: u32) -> Extend<'_> { .extend(views[start..start + len].iter().map(|v| { // A view of 12 bytes or fewer stores the string inline and // has no buffer index, so adding 0 leaves it unchanged. - let offset = if *v as u32 > 12 { buffer_offset } else { 0 }; + let offset = if *v as u32 > MAX_INLINE_VIEW_LEN { + buffer_offset + } else { + 0 + }; // The buffer index is bits 64 to 95 of a view. The new index // is less than the number of data buffers, so the sum never // carries into the bits above it. From 5352e8af9f9a7dd951717d7ad6f1a62b9861d367 Mon Sep 17 00:00:00 2001 From: Neil Conway Date: Wed, 21 Oct 2026 09:40:40 -0400 Subject: [PATCH 5/5] chore: clarify the comments on the string view buffer offset The comment on the shifted add said that "the new index is less than the number of data buffers", without saying which index or why that rules out a carry. It now says that adding `buffer_offset` gives the string's buffer index in the output, which fits in 32 bits: `MutableArrayData::with_capacities` panics if the total number of data buffers in its input arrays doesn't fit in 32 bits. The comment on the length check now speaks of a view of a string of 12 bytes or fewer, since every view is 16 bytes. --- arrow-data/src/transform/mod.rs | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/arrow-data/src/transform/mod.rs b/arrow-data/src/transform/mod.rs index f72d781e9546..5937d9fa3211 100644 --- a/arrow-data/src/transform/mod.rs +++ b/arrow-data/src/transform/mod.rs @@ -226,16 +226,18 @@ fn build_extend_view(array: &ArrayData, buffer_offset: u32) -> Extend<'_> { mutable .buffer1 .extend(views[start..start + len].iter().map(|v| { - // A view of 12 bytes or fewer stores the string inline and - // has no buffer index, so adding 0 leaves it unchanged. + // A view of a string of 12 bytes or fewer stores the + // string inline and has no buffer index, so adding 0 + // leaves it unchanged. let offset = if *v as u32 > MAX_INLINE_VIEW_LEN { buffer_offset } else { 0 }; - // The buffer index is bits 64 to 95 of a view. The new index - // is less than the number of data buffers, so the sum never - // carries into the bits above it. + // The buffer index is bits 64 to 95 of a view. Adding + // `buffer_offset` gives the string's buffer index in the + // output, which fits in 32 bits, so the sum never carries + // into bit 96. v + ((offset as u128) << 64) })); Ok(())