bench: add MutableBuffer::extend, enhance MutableArrayData

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

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);

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);