benchmark/perf: add wide column take benchmarks

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<ArrayRef> {
+    types
+        .iter()
+        .map(|dt| create_array_for_type(dt, size, null_density))
+        .collect()
+}
+
+fn make_record_batch(columns: Vec<ArrayRef>) -> 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);

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::<Int8Type>(size, null_density)),
+        DataType::Int16 => Arc::new(create_primitive_array::<Int16Type>(size, null_density)),
+        DataType::Int32 => Arc::new(create_primitive_array::<Int32Type>(size, null_density)),
+        DataType::Int64 => Arc::new(create_primitive_array::<Int64Type>(size, null_density)),
+        DataType::UInt8 => Arc::new(create_primitive_array::<UInt8Type>(size, null_density)),
+        DataType::UInt16 => Arc::new(create_primitive_array::<UInt16Type>(size, null_density)),
+        DataType::UInt32 => Arc::new(create_primitive_array::<UInt32Type>(size, null_density)),
+        DataType::UInt64 => Arc::new(create_primitive_array::<UInt64Type>(size, null_density)),
+        DataType::Float32 => Arc::new(create_primitive_array::<Float32Type>(size, null_density)),
+        DataType::Float64 => Arc::new(create_primitive_array::<Float64Type>(size, null_density)),
+        DataType::Utf8 => Arc::new(create_string_array::<i32>(size, null_density)),
+        DataType::LargeUtf8 => Arc::new(create_string_array::<i64>(size, null_density)),
+        DataType::Utf8View => Arc::new(create_string_view_array(size, null_density)),
+        DataType::Binary => Arc::new(create_binary_array::<i32>(size, null_density)),
+        DataType::LargeBinary => Arc::new(create_binary_array::<i64>(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}"),
+    }
+}