extern crate criterion;
use arrow::{array::PrimitiveArray, datatypes::Int64Type};
use criterion::{black_box, criterion_group, criterion_main, Criterion};
use datafusion_expr::{ColumnarValue, ScalarFunctionArgs};
use datafusion_functions::string::chr;
use rand::{Rng, SeedableRng};
use arrow::datatypes::{DataType, Field};
use rand::rngs::StdRng;
use std::sync::Arc;
pub fn seedable_rng() -> StdRng {
StdRng::seed_from_u64(42)
}
fn criterion_benchmark(c: &mut Criterion) {
let cot_fn = chr();
let size = 1024;
let input: PrimitiveArray<Int64Type> = {
let null_density = 0.2;
let mut rng = StdRng::seed_from_u64(42);
(0..size)
.map(|_| {
if rng.random::<f32>() < null_density {
None
} else {
Some(rng.random_range::<i64, _>(1i64..10_000))
}
})
.collect()
};
let input = Arc::new(input);
let args = vec![ColumnarValue::Array(input)];
let arg_fields = args
.iter()
.enumerate()
.map(|(idx, arg)| Field::new(format!("arg_{idx}"), arg.data_type(), true).into())
.collect::<Vec<_>>();
c.bench_function("chr", |b| {
b.iter(|| {
black_box(
cot_fn
.invoke_with_args(ScalarFunctionArgs {
args: args.clone(),
arg_fields: arg_fields.clone(),
number_rows: size,
return_field: Field::new("f", DataType::Utf8, true).into(),
})
.unwrap(),
)
})
});
}
criterion_group!(benches, criterion_benchmark);
criterion_main!(benches);