use rand::{rngs::StdRng, Rng, SeedableRng};
use criterion::{criterion_group, criterion_main, Criterion};
use arrow2::compute::take;
use arrow2::util::bench_util::*;
use arrow2::{array::*, datatypes::DataType};
pub fn seedable_rng() -> StdRng {
StdRng::seed_from_u64(42)
}
fn create_random_index(size: usize, null_density: f32) -> PrimitiveArray<i32> {
let mut rng = seedable_rng();
(0..size)
.map(|_| {
if rng.gen::<f32>() > null_density {
let value = rng.gen_range::<i32, _>(0i32..size as i32);
Some(value)
} else {
None
}
})
.collect::<PrimitiveArray<i32>>()
}
fn bench_take(values: &dyn Array, indices: &PrimitiveArray<i32>) {
criterion::black_box(take::take(values, indices).unwrap());
}
fn add_benchmark(c: &mut Criterion) {
(10..=20).step_by(2).for_each(|log2_size| {
let size = 2usize.pow(log2_size);
let values = create_primitive_array::<i32>(size, DataType::Int32, 0.0);
let values_nulls = create_primitive_array::<i32>(size, DataType::Int32, 0.2);
let indices = create_random_index(size, 0.0);
let indices_nulls = create_random_index(size, 0.5);
c.bench_function(&format!("take i32 2^{}", log2_size), |b| {
b.iter(|| bench_take(&values, &indices))
});
c.bench_function(&format!("take i32 nulls 2^{}", log2_size), |b| {
b.iter(|| bench_take(&values, &indices_nulls))
});
c.bench_function(&format!("take i32 values nulls 2^{}", log2_size), |b| {
b.iter(|| bench_take(&values_nulls, &indices))
});
let values = create_boolean_array(size, 0.0, 0.5);
c.bench_function(&format!("take bool 2^{}", log2_size), |b| {
b.iter(|| bench_take(&values, &indices))
});
c.bench_function(&format!("take bool nulls 2^{}", log2_size), |b| {
b.iter(|| bench_take(&values, &indices_nulls))
});
let values = create_string_array::<i32>(512, 4, 0.0, 42);
c.bench_function(&format!("take str 2^{}", log2_size), |b| {
b.iter(|| bench_take(&values, &indices))
});
let values = create_string_array::<i32>(512, 4, 0.0, 42);
c.bench_function(&format!("take str nulls 2^{}", log2_size), |b| {
b.iter(|| bench_take(&values, &indices_nulls))
});
});
let values = create_string_array::<i32>(512, 4, 0.0, 42);
let indices = create_random_index(512, 0.5);
c.bench_function("take str null indices 512", |b| {
b.iter(|| bench_take(&values, &indices))
});
let values = create_string_array::<i32>(1024, 4, 0.0, 42);
let indices = create_random_index(1024, 0.5);
c.bench_function("take str null indices 1024", |b| {
b.iter(|| bench_take(&values, &indices))
});
let values = create_string_array::<i32>(1024, 4, 0.5, 42);
let indices = create_random_index(1024, 0.0);
c.bench_function("take str null values 1024", |b| {
b.iter(|| bench_take(&values, &indices))
});
let values = create_string_array::<i32>(1024, 4, 0.5, 42);
let indices = create_random_index(1024, 0.5);
c.bench_function("take str null values null indices 1024", |b| {
b.iter(|| bench_take(&values, &indices))
});
}
criterion_group!(benches, add_benchmark);
criterion_main!(benches);