diff --git a/datafusion/functions/benches/unicode_expressions/left_right.rs b/datafusion/functions/benches/unicode_expressions/left_right.rs index 01a0dcad6990b..7044df29367c4 100644 --- a/datafusion/functions/benches/unicode_expressions/left_right.rs +++ b/datafusion/functions/benches/unicode_expressions/left_right.rs @@ -16,55 +16,78 @@ // under the License. use std::hint::black_box; -use std::ops::Range; +use std::ops::{Range, RangeInclusive}; use std::sync::Arc; -use arrow::array::{ArrayRef, Int64Array}; -use arrow::datatypes::{DataType, Field}; +use arrow::array::StringViewArray; +use arrow::datatypes::{Field, Int64Type}; use arrow::util::bench_util::{ - create_string_array_with_len, create_string_view_array_with_len, + create_primitive_array_range, create_string_array_with_len_range_and_prefix_and_seed, }; use criterion::{Criterion, criterion_group}; +use datafusion_common::ScalarValue; use datafusion_common::config::ConfigOptions; -use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; +use datafusion_expr::{ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs}; use datafusion_functions::unicode::{left, right}; const BATCH_SIZE: usize = 8192; +/// How the `n` argument is passed. +enum NArg { + /// The same `n` for every row, as in `left(s, 5)`. + Scalar(i64), + /// A different `n` for each row, drawn from the range. + PerRow(Range), +} + fn create_args( - str_len: usize, - n_range: Range, + str_len: &RangeInclusive, + n: &NArg, is_string_view: bool, ) -> Vec { + let strings = create_string_array_with_len_range_and_prefix_and_seed::( + BATCH_SIZE, + 0.1, + *str_len.start(), + *str_len.end(), + "", + 42, + ); let string_arg = if is_string_view { - ColumnarValue::Array(Arc::new(create_string_view_array_with_len( - BATCH_SIZE, 0.1, str_len, true, - ))) + ColumnarValue::Array(Arc::new(strings.iter().collect::())) } else { - ColumnarValue::Array(Arc::new(create_string_array_with_len::( - BATCH_SIZE, 0.1, str_len, - ))) + ColumnarValue::Array(Arc::new(strings)) }; - let n_span = (n_range.end - n_range.start) as usize; - let n_values: Vec = (0..BATCH_SIZE) - .map(|i| n_range.start + (i % n_span) as i64) - .collect(); - let n_array = Arc::new(Int64Array::from(n_values)); + let n_arg = match n { + NArg::Scalar(n) => ColumnarValue::Scalar(ScalarValue::Int64(Some(*n))), + NArg::PerRow(range) => ColumnarValue::Array(Arc::new( + create_primitive_array_range::(BATCH_SIZE, 0.0, range.clone()), + )), + }; - vec![ - string_arg, - ColumnarValue::Array(Arc::clone(&n_array) as ArrayRef), - ] + vec![string_arg, n_arg] } fn criterion_benchmark(c: &mut Criterion) { - // Short results (1-10 chars) produce inline StringView entries (≤12 bytes). - // Long results (20-29 chars) produce out-of-line entries. + // Input lengths vary within each case, as in real data: with fixed-length + // inputs, per-row work that depends on the input length is unrealistically + // predictable. let cases = [ - ("short_result", 32, 1..11_i64), - ("long_result", 32, 20..30_i64), + // Results of up to 5 chars, stored inline in StringView arrays (≤12 bytes). + ("short_result", 1..=32, NArg::Scalar(5)), + // 25-char results, stored out of line. + ("long_result", 32..=64, NArg::Scalar(25)), + // Short results from long inputs. + ("short_result_long_input", 96..=256, NArg::Scalar(5)), + // `n` exceeds every input's length, so each result is the whole input. + ("n_exceeds_len", 1..=32, NArg::Scalar(40)), + // Negative `n` removes characters from the other end. + ("negative_n", 1..=32, NArg::Scalar(-5)), + // `n` computed per row, as in `left(s, strpos(s, '-') - 1)`. + ("per_row_n", 1..=32, NArg::PerRow(1..11)), ]; + let config_options = Arc::new(ConfigOptions::default()); for function in [left(), right()] { let mut group = c.benchmark_group(function.name().to_string()); @@ -76,9 +99,9 @@ fn criterion_benchmark(c: &mut Criterion) { "string" }; - for (case_name, str_len, n_range) in &cases { + for (case_name, str_len, n) in &cases { let bench_name = format!("{array_type} {case_name}"); - let args = create_args(*str_len, n_range.clone(), is_string_view); + let args = create_args(str_len, n, is_string_view); let arg_fields: Vec<_> = args .iter() .enumerate() @@ -86,8 +109,13 @@ fn criterion_benchmark(c: &mut Criterion) { Field::new(format!("arg_{idx}"), arg.data_type(), true).into() }) .collect(); - let config_options = Arc::new(ConfigOptions::default()); - let return_field = Field::new("f", DataType::Utf8View, true).into(); + let scalar_arguments = vec![None; arg_fields.len()]; + let return_field = function + .return_field_from_args(ReturnFieldArgs { + arg_fields: &arg_fields, + scalar_arguments: &scalar_arguments, + }) + .expect("should resolve return field"); group.bench_function(&bench_name, |b| { b.iter(|| {