diff --git a/arrow/benches/take_kernels.rs b/arrow/benches/take_kernels.rs index a10f80c5907b..5eb9d4fde1af 100644 --- a/arrow/benches/take_kernels.rs +++ b/arrow/benches/take_kernels.rs @@ -23,11 +23,13 @@ use rand::Rng; extern crate arrow; -use arrow::compute::{TakeOptions, take}; +use arrow::compute::{TakeOptions, take, take_record_batch}; use arrow::datatypes::*; +use arrow::record_batch::RecordBatch; use arrow::util::test_util::seedable_rng; use arrow::{array::*, util::bench_util::*}; use std::hint; +use std::sync::Arc; fn create_random_index(size: usize, null_density: f32) -> UInt32Array { let mut rng = seedable_rng(); @@ -47,6 +49,26 @@ fn bench_take(values: &dyn Array, indices: &UInt32Array) { hint::black_box(take(values, indices, None).unwrap()); } +fn create_columns(types: &[DataType], size: usize, null_density: f32) -> Vec { + types + .iter() + .map(|dt| create_array_for_type(dt, size, null_density)) + .collect() +} + +fn make_record_batch(columns: Vec) -> RecordBatch { + let fields: Vec<_> = columns + .iter() + .enumerate() + .map(|(i, col)| Field::new(format!("c{i}"), col.data_type().clone(), true)) + .collect(); + RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap() +} + +fn bench_take_record_batch(batch: &RecordBatch, indices: &UInt32Array) { + hint::black_box(take_record_batch(batch, indices).unwrap()); +} + fn bench_take_bounds_check(values: &dyn Array, indices: &UInt32Array) { hint::black_box(take(values, indices, Some(TakeOptions { check_bounds: true })).unwrap()); } @@ -277,6 +299,65 @@ fn add_benchmark(c: &mut Criterion) { "take fsb value optimized len: 16, null values, indices: 1024", |b| b.iter(|| bench_take(&values, &indices)), ); + + let types = [ + DataType::Int32, + DataType::Int64, + DataType::Float32, + DataType::Float64, + DataType::Boolean, + ]; + let batch = make_record_batch(create_columns(&types, 1024, 0.0)); + let indices = create_random_index(1024, 0.0); + c.bench_function("take_record_batch 5 primitive cols no nulls 1024", |b| { + b.iter(|| bench_take_record_batch(&batch, &indices)) + }); + + let types = [ + DataType::Utf8, + DataType::LargeUtf8, + DataType::Utf8View, + DataType::Binary, + DataType::LargeBinary, + DataType::FixedSizeBinary(16), + ]; + let batch = make_record_batch(create_columns(&types, 1024, 0.0)); + let indices = create_random_index(1024, 0.0); + c.bench_function( + "take_record_batch 6 string/binary cols no nulls 1024", + |b| b.iter(|| bench_take_record_batch(&batch, &indices)), + ); + + let types = [ + DataType::Int32, + DataType::Utf8, + DataType::Float64, + DataType::Boolean, + DataType::Utf8View, + DataType::Int64, + DataType::Binary, + ]; + let batch = make_record_batch(create_columns(&types, 1024, 0.5)); + let indices = create_random_index(1024, 0.0); + c.bench_function("take_record_batch 7 mixed cols null values 1024", |b| { + b.iter(|| bench_take_record_batch(&batch, &indices)) + }); + + let types = [ + DataType::Int32, + DataType::Utf8, + DataType::Float64, + DataType::Boolean, + DataType::Utf8View, + DataType::Int64, + DataType::Binary, + ]; + let batch = make_record_batch(create_columns(&types, 1024, 0.5)); + let indices = create_random_index(1024, 0.5); + c.bench_function( + "take_record_batch 7 mixed cols null values null indices 1024", + |b| b.iter(|| bench_take_record_batch(&batch, &indices)), + ); } criterion_group!(benches, add_benchmark); diff --git a/arrow/src/util/bench_util.rs b/arrow/src/util/bench_util.rs index cde0bc1d20f8..4edd0e50bc62 100644 --- a/arrow/src/util/bench_util.rs +++ b/arrow/src/util/bench_util.rs @@ -947,3 +947,30 @@ where } builder.finish() } + +/// Creates a random array for the given [`DataType`], `size`, and `null_density`. +/// +/// Useful for building arrays and record batches in benchmarks without +/// repeating per-type construction logic. Panics on unsupported types. +pub fn create_array_for_type(data_type: &DataType, size: usize, null_density: f32) -> ArrayRef { + match data_type { + DataType::Boolean => Arc::new(create_boolean_array(size, null_density, 0.5)), + DataType::Int8 => Arc::new(create_primitive_array::(size, null_density)), + DataType::Int16 => Arc::new(create_primitive_array::(size, null_density)), + DataType::Int32 => Arc::new(create_primitive_array::(size, null_density)), + DataType::Int64 => Arc::new(create_primitive_array::(size, null_density)), + DataType::UInt8 => Arc::new(create_primitive_array::(size, null_density)), + DataType::UInt16 => Arc::new(create_primitive_array::(size, null_density)), + DataType::UInt32 => Arc::new(create_primitive_array::(size, null_density)), + DataType::UInt64 => Arc::new(create_primitive_array::(size, null_density)), + DataType::Float32 => Arc::new(create_primitive_array::(size, null_density)), + DataType::Float64 => Arc::new(create_primitive_array::(size, null_density)), + DataType::Utf8 => Arc::new(create_string_array::(size, null_density)), + DataType::LargeUtf8 => Arc::new(create_string_array::(size, null_density)), + DataType::Utf8View => Arc::new(create_string_view_array(size, null_density)), + DataType::Binary => Arc::new(create_binary_array::(size, null_density)), + DataType::LargeBinary => Arc::new(create_binary_array::(size, null_density)), + DataType::FixedSizeBinary(n) => Arc::new(create_fsb_array(size, null_density, *n as usize)), + other => panic!("unsupported data type for create_array_for_type: {other}"), + } +}