bench: Use variable-length inputs and scalar n in left/right benchmark

The Utf8 inputs were all exactly 32 bytes, while the Utf8View inputs
used a different length distribution, so the two array types were not
benchmarked on the same data. Fixed-length inputs also make per-row
work that depends on the input length, such as an ASCII check over the
whole string, perfectly predictable, so it looks nearly free. And every
case passed `n` as an array cycling through a short range, although `n`
is usually a literal.

Generate variable-length inputs for every case and build the Utf8 and
Utf8View arrays from the same strings. Pass `n` as a scalar, and add a
case with a different random `n` per row. Add cases for short results
from long inputs, for `n` exceeding the input length, and for negative
`n`.

Also derive the return field from the function instead of always using
Utf8View. Since #23330, `left` and `right` return Utf8 for Utf8 input,
and the mismatch fails the check in `ScalarUDF::invoke_with_args` when
debug assertions are enabled.

Author: neilconway

datafusion/functions/benches/unicode_expressions/left_right.rs

@@ -16,55 +16,82 @@
 // under the License.
 
 use std::hint::black_box;
-use std::ops::Range;
+use std::ops::RangeInclusive;
 use std::sync::Arc;
 
-use arrow::array::{ArrayRef, Int64Array};
-use arrow::datatypes::{DataType, Field};
-use arrow::util::bench_util::{
-    create_string_array_with_len, create_string_view_array_with_len,
-};
+use arrow::array::{Int64Array, StringViewArray};
+use arrow::datatypes::Field;
+use arrow::util::bench_util::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};
+use rand::rngs::StdRng;
+use rand::{Rng, SeedableRng};
 
 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(RangeInclusive<i64>),
+}
+
 fn create_args(
-    str_len: usize,
-    n_range: Range<i64>,
+    str_len: RangeInclusive<usize>,
+    n: &NArg,
     is_string_view: bool,
 ) -> Vec<ColumnarValue> {
+    let strings = create_string_array_with_len_range_and_prefix_and_seed::<i32>(
+        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::<StringViewArray>()))
     } else {
-        ColumnarValue::Array(Arc::new(create_string_array_with_len::<i32>(
-            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<i64> = (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) => {
+            let mut rng = StdRng::seed_from_u64(42);
+            let n_values = Int64Array::from_iter_values(
+                (0..BATCH_SIZE).map(|_| rng.random_range(range.clone())),
+            );
+            ColumnarValue::Array(Arc::new(n_values))
+        }
+    };
 
-    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..=10)),
     ];
+    let config_options = Arc::new(ConfigOptions::default());
 
     for function in [left(), right()] {
         let mut group = c.benchmark_group(function.name().to_string());
@@ -76,18 +103,29 @@ 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.clone(), n, is_string_view);
                 let arg_fields: Vec<_> = args
                     .iter()
                     .enumerate()
                     .map(|(idx, arg)| {
                         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<_> = args
+                    .iter()
+                    .map(|arg| match arg {
+                        ColumnarValue::Scalar(scalar) => Some(scalar),
+                        ColumnarValue::Array(_) => None,
+                    })
+                    .collect();
+                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(|| {