extern crate criterion;
use arrow::{
datatypes::{Field, Float32Type, Float64Type},
util::bench_util::create_primitive_array,
};
use criterion::{black_box, criterion_group, criterion_main, Criterion};
use datafusion_expr::{ColumnarValue, ScalarFunctionArgs};
use datafusion_functions::math::trunc;
use arrow::datatypes::DataType;
use std::sync::Arc;
fn criterion_benchmark(c: &mut Criterion) {
let trunc = trunc();
for size in [1024, 4096, 8192] {
let f32_array = Arc::new(create_primitive_array::<Float32Type>(size, 0.2));
let f32_args = vec![ColumnarValue::Array(f32_array)];
let arg_fields = vec![Field::new("a", DataType::Float32, false).into()];
let return_field = Field::new("f", DataType::Float32, true).into();
c.bench_function(&format!("trunc f32 array: {size}"), |b| {
b.iter(|| {
black_box(
trunc
.invoke_with_args(ScalarFunctionArgs {
args: f32_args.clone(),
arg_fields: arg_fields.clone(),
number_rows: size,
return_field: Arc::clone(&return_field),
})
.unwrap(),
)
})
});
let f64_array = Arc::new(create_primitive_array::<Float64Type>(size, 0.2));
let f64_args = vec![ColumnarValue::Array(f64_array)];
let arg_fields = vec![Field::new("a", DataType::Float64, true).into()];
let return_field = Field::new("f", DataType::Float64, true).into();
c.bench_function(&format!("trunc f64 array: {size}"), |b| {
b.iter(|| {
black_box(
trunc
.invoke_with_args(ScalarFunctionArgs {
args: f64_args.clone(),
arg_fields: arg_fields.clone(),
number_rows: size,
return_field: Arc::clone(&return_field),
})
.unwrap(),
)
})
});
}
}
criterion_group!(benches, criterion_benchmark);
criterion_main!(benches);