extern crate criterion;
use arrow::{
datatypes::{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::cot;
use arrow::datatypes::DataType;
use std::sync::Arc;
fn criterion_benchmark(c: &mut Criterion) {
let cot_fn = cot();
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)];
c.bench_function(&format!("cot f32 array: {}", size), |b| {
b.iter(|| {
black_box(
cot_fn
.invoke_with_args(ScalarFunctionArgs {
args: f32_args.clone(),
number_rows: size,
return_type: &DataType::Float32,
})
.unwrap(),
)
})
});
let f64_array = Arc::new(create_primitive_array::<Float64Type>(size, 0.2));
let f64_args = vec![ColumnarValue::Array(f64_array)];
c.bench_function(&format!("cot f64 array: {}", size), |b| {
b.iter(|| {
black_box(
cot_fn
.invoke_with_args(ScalarFunctionArgs {
args: f64_args.clone(),
number_rows: size,
return_type: &DataType::Float64,
})
.unwrap(),
)
})
});
}
}
criterion_group!(benches, criterion_benchmark);
criterion_main!(benches);