This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new 25569a3328 benchmark/perf: add wide column take benchmarks (#10442)
25569a3328 is described below
commit 25569a332896a0385638e03e11630a37f18c22fa
Author: RIchard Baah <[email protected]>
AuthorDate: Mon Jul 27 12:00:23 2026 -0400
benchmark/perf: add wide column take benchmarks (#10442)
# Which issue does this PR close?
<!--
We generally require a GitHub issue to be filed for all bug fixes and
enhancements and this helps us generate change logs for our releases.
You can link an issue to this PR using the GitHub syntax.
-->
- works towards #8879.
# Rationale for this change
see
https://github.com/apache/arrow-rs/issues/8879#issuecomment-5086286171
the point of these benchmarks are to measure the change of #10441
<!--
Why are you proposing this change? If this is already explained clearly
in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand
your changes and offer better suggestions for fixes.
-->
# What changes are included in this PR?
adds benchmarks for `take_record_batch()` .
<!--
There is no need to duplicate the description in the issue here but it
is sometimes worth providing a summary of the individual changes in this
PR.
-->
# Are these changes tested?
n/a
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example,
are they covered by existing tests)?
If this PR claims a performance improvement, please include evidence
such as benchmark results.
-->
# Are there any user-facing changes?
no
<!--
If there are user-facing changes then we may require documentation to be
updated before approving the PR.
If there are any breaking changes to public APIs, please call them out.
-->
---
arrow/benches/take_kernels.rs | 83 ++++++++++++++++++++++++++++++++++++++++++-
arrow/src/util/bench_util.rs | 27 ++++++++++++++
2 files changed, 109 insertions(+), 1 deletion(-)
diff --git a/arrow/benches/take_kernels.rs b/arrow/benches/take_kernels.rs
index a10f80c590..5eb9d4fde1 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<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);
diff --git a/arrow/src/util/bench_util.rs b/arrow/src/util/bench_util.rs
index cde0bc1d20..4edd0e50bc 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::<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}"),
+ }
+}
