diff --git a/arrow/benches/interleave_byte_view_cases/mod.rs b/arrow/benches/interleave_byte_view_cases/mod.rs new file mode 100644 index 000000000000..0f33961b720c --- /dev/null +++ b/arrow/benches/interleave_byte_view_cases/mod.rs @@ -0,0 +1,197 @@ +// 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. + +use arrow::array::{StringViewArray, StringViewBuilder}; +use arrow_buffer::Buffer; +use rand::{RngExt, SeedableRng, rngs::StdRng}; + +const SOURCE_COUNT: usize = 4; +const ROWS_PER_SOURCE: usize = 8192; +const ALIASED_VALUES: usize = 64; + +pub struct ByteViewCase { + pub name: String, + pub arrays: Vec, + pub indices: Vec<(usize, usize)>, +} + +#[derive(Clone, Copy)] +enum Layout { + Unique, + Random, + Sparse, + RepeatedIndices, + SharedRanges, + SlicedNulls, +} + +impl Layout { + fn name(self) -> &'static str { + match self { + Self::Unique => "unique", + Self::Random => "random", + Self::Sparse => "sparse", + Self::RepeatedIndices => "repeated_indices", + Self::SharedRanges => "shared_ranges", + Self::SlicedNulls => "sliced_nulls", + } + } +} + +fn value(id: usize, min_len: usize, max_len: usize) -> String { + let len = min_len + id.wrapping_mul(4051) % (max_len - min_len + 1); + // Include the row identity even in the shortest values, so the unique + // controls do not accidentally contain equal strings. + let mut value = format!("{id:08x}:"); + value.extend((value.len()..len).map(|i| (b'a' + ((id + i * 17) % 26) as u8) as char)); + value +} + +fn make_case(layout: Layout, selected: usize, min_len: usize, max_len: usize) -> ByteViewCase { + let arrays = if matches!(layout, Layout::SharedRanges) { + let mut bytes = Vec::new(); + let mut ranges = Vec::new(); + for id in 0..ALIASED_VALUES { + let value = value(id, min_len, max_len); + ranges.push((bytes.len() as u32, value.len() as u32)); + bytes.extend_from_slice(value.as_bytes()); + } + let buffer = Buffer::from_vec(bytes); + (0..SOURCE_COUNT) + .map(|source| { + let mut builder = StringViewBuilder::with_capacity(ROWS_PER_SOURCE); + let block = builder.append_block(buffer.clone()); + for row in 0..ROWS_PER_SOURCE { + let (offset, len) = ranges[(row + source * 17) % ALIASED_VALUES]; + builder.try_append_view(block, offset, len).unwrap(); + } + builder.finish() + }) + .collect() + } else { + (0..SOURCE_COUNT) + .map(|source| { + let sliced = matches!(layout, Layout::SlicedNulls); + let extra = if sliced { 37 } else { 0 }; + let block_size = if matches!(layout, Layout::Sparse) { + 1024 + } else { + 64 * 1024 + }; + let mut builder = StringViewBuilder::with_capacity(ROWS_PER_SOURCE + extra) + .with_fixed_block_size(block_size); + for row in 0..ROWS_PER_SOURCE + extra { + if sliced && row % 7 == 0 { + builder.append_null(); + } else if sliced && row % 11 == 0 { + builder.append_value("inline"); + } else { + builder.append_value(value( + source * (ROWS_PER_SOURCE + extra) + row, + min_len, + max_len, + )); + } + } + let array = builder.finish(); + if sliced { + array.slice(17, ROWS_PER_SOURCE) + } else { + array + } + }) + .collect() + }; + let mut rng = StdRng::seed_from_u64(0); + let indices = (0..selected) + .map(|i| { + if matches!(layout, Layout::Random) { + return ( + rng.random_range(0..SOURCE_COUNT), + rng.random_range(0..ROWS_PER_SOURCE), + ); + } + let source = i % SOURCE_COUNT; + let row = i / SOURCE_COUNT; + let row = if matches!(layout, Layout::RepeatedIndices) { + row % 16 + } else { + row + }; + // An odd multiplier permutes the power-of-two source length, + // distributing selections across its buffers without replacement. + (source, (row * 4051) % ROWS_PER_SOURCE) + }) + .collect(); + ByteViewCase { + name: format!( + "{}_len{min_len}-{max_len}_selected{selected}", + layout.name() + ), + arrays, + indices, + } +} + +fn source_buffer_case(buffer_count: usize, selected: usize) -> ByteViewCase { + let rows = SOURCE_COUNT * ROWS_PER_SOURCE / 2; + let buffers_per_source = buffer_count / 2; + let block_size = (rows / buffers_per_source * 16) as u32; + let arrays: Vec<_> = (0..2) + .map(|source| { + let mut builder = + StringViewBuilder::with_capacity(rows).with_fixed_block_size(block_size); + for row in 0..rows { + builder.append_value(value(source * rows + row, 16, 16)); + } + let array = builder.finish(); + assert_eq!(array.data_buffers().len(), buffers_per_source); + array + }) + .collect(); + ByteViewCase { + name: format!("source_buffers{buffer_count}_selected{selected}"), + arrays, + indices: (0..selected) + .map(|i| (i % 2, ((i / 2) * 4051) % rows)) + .collect(), + } +} + +pub fn for_each_case(mut visit: impl FnMut(ByteViewCase)) { + for (min_len, max_len) in [(13, 20), (20, 60), (100, 400)] { + for selected in [512, 8192] { + for layout in [ + Layout::Unique, + Layout::Random, + Layout::Sparse, + Layout::RepeatedIndices, + Layout::SharedRanges, + Layout::SlicedNulls, + ] { + visit(make_case(layout, selected, min_len, max_len)); + } + } + } + // Keep the same source rows, payload, and selections while independently + // varying the number of source buffers that interleave must track. + for buffers in [4, 4096, 32768] { + for selected in [3, 32] { + visit(source_buffer_case(buffers, selected)); + } + } +} diff --git a/arrow/benches/interleave_kernels.rs b/arrow/benches/interleave_kernels.rs index ee5bfb13d980..e5e52d844bd7 100644 --- a/arrow/benches/interleave_kernels.rs +++ b/arrow/benches/interleave_kernels.rs @@ -18,7 +18,9 @@ #[macro_use] extern crate criterion; -use criterion::Criterion; +mod interleave_byte_view_cases; + +use criterion::{BenchmarkId, Criterion, Throughput}; use std::ops::Range; use rand::RngExt; @@ -30,6 +32,7 @@ use arrow_buffer::ScalarBuffer; use arrow_select::interleave::interleave; use std::hint; use std::sync::Arc; +use std::time::Duration; fn do_bench( c: &mut Criterion, @@ -66,6 +69,31 @@ fn bench_values(c: &mut Criterion, name: &str, len: usize, values: &[&dyn Array] }); } +fn bench_byte_view_compaction(c: &mut Criterion) { + let mut group = c.benchmark_group("interleave_byte_view_compaction"); + interleave_byte_view_cases::for_each_case(|case| { + let values: Vec<&dyn Array> = case.arrays.iter().map(|a| a as &dyn Array).collect(); + group.throughput(Throughput::Elements(case.indices.len() as u64)); + group.bench_with_input( + BenchmarkId::new("interleave", &case.name), + &case.indices, + |b, indices| b.iter(|| hint::black_box(interleave(&values, indices).unwrap())), + ); + group.bench_with_input( + BenchmarkId::new("interleave_gc", &case.name), + &case.indices, + |b, indices| { + b.iter(|| { + let selected = interleave(&values, indices).unwrap(); + let selected = selected.as_any().downcast_ref::().unwrap(); + hint::black_box(Arc::new(selected.gc()) as ArrayRef) + }) + }, + ); + }); + group.finish(); +} + fn add_benchmark(c: &mut Criterion) { let i32 = create_primitive_array::(1024, 0.); let i32_opt = create_primitive_array::(1024, 0.5); @@ -219,4 +247,12 @@ fn add_benchmark(c: &mut Criterion) { } criterion_group!(benches, add_benchmark); -criterion_main!(benches); +criterion_group! { + name = byte_view_benches; + config = Criterion::default() + .sample_size(20) + .warm_up_time(Duration::from_millis(100)) + .measurement_time(Duration::from_millis(300)); + targets = bench_byte_view_compaction +} +criterion_main!(benches, byte_view_benches);