perf: Optimize substr for Utf8, LargeUtf8 (#21366)

## Which issue does this PR close?

- Closes #21364.

## Rationale for this change

For `Utf8` and `LargeUtf8` inputs, we can optimize `substr` to avoid
copying the output strings; instead, we can return a `StringViewArray`
that points into the input value buffer.

Benchmarks (M4 Max):

no count, short strings (size=1024):
    - string_view:  5.97 µs -> 5.96 µs (-0.2%)
    - string:       7.80 µs -> 4.99 µs (-36.1%)
    - large_string: 8.47 µs -> 4.90 µs (-42.2%)

  no count, short strings (size=4096):
    - string_view:  23.10 µs -> 22.90 µs (-0.9%)
    - string:       31.24 µs -> 18.31 µs (-41.4%)
    - large_string: 34.10 µs -> 17.70 µs (-48.1%)

  with count, long strings (size=1024, count=64, strlen=128):
    - string_view:  10.16 µs -> 10.79 µs (+6.2%)
    - string:       11.90 µs -> 8.38 µs (-29.6%)
    - large_string: 11.93 µs -> 8.30 µs (-30.5%)

  with count, long strings (size=4096, count=64, strlen=128):
    - string_view:  39.37 µs -> 38.79 µs (-1.5%)
    - string:       46.22 µs -> 30.25 µs (-34.6%)
    - large_string: 46.57 µs -> 30.49 µs (-34.5%)

  short count, long strings (size=1024, count=6, strlen=128):
    - string_view:  11.65 µs -> 11.57 µs (-0.7%)
    - string:       14.97 µs -> 11.37 µs (-24.1%)
    - large_string: 14.92 µs -> 11.37 µs (-23.8%)

  short count, long strings (size=4096, count=6, strlen=128):
    - string_view:  45.88 µs -> 43.82 µs (-4.5%)
    - string:       58.38 µs -> 43.55 µs (-25.4%)
    - large_string: 58.59 µs -> 43.58 µs (-25.6%)

  scalar start, no count, short strings (size=1024, strlen=12):
    - string_view:  6.07 µs -> 6.10 µs (+0.5%)
    - string:       7.81 µs -> 5.06 µs (-35.2%)

  scalar start, no count, short strings (size=4096, strlen=12):
    - string_view:  23.08 µs -> 22.62 µs (-2.0%)
    - string:       31.07 µs -> 18.86 µs (-39.3%)

  scalar start, no count, long strings (size=1024, strlen=128):
    - string_view:  9.99 µs -> 10.65 µs (+6.6%)
    - string:       12.01 µs -> 8.17 µs (-32.0%)

  scalar start, no count, long strings (size=4096, strlen=128):
    - string_view:  38.57 µs -> 39.79 µs (+3.2%)
    - string:       46.83 µs -> 31.67 µs (-32.4%)

  scalar start=1, no count, long strings (size=1024, strlen=128):
    - string_view:  9.78 µs -> 10.48 µs (+7.2%)
    - string:       12.02 µs -> 8.16 µs (-32.1%)

  scalar start=1, no count, long strings (size=4096, strlen=128):
    - string_view:  38.54 µs -> 40.18 µs (+4.3%)
    - string:       46.36 µs -> 31.73 µs (-31.6%)

  scalar args, short strings (size=1024, count=6, strlen=12):
    - string_view:  11.30 µs -> 11.23 µs (-0.7%)
    - string:       15.04 µs -> 11.52 µs (-23.4%)

  scalar args, short strings (size=4096, count=6, strlen=12):
    - string_view:  44.34 µs -> 43.98 µs (-0.8%)
    - string:       59.63 µs -> 45.02 µs (-24.5%)

  scalar args, long strings (size=1024, count=64, strlen=128):
    - string_view:  10.51 µs -> 12.05 µs (+14.6%)
    - string:       12.21 µs -> 8.67 µs (-28.9%)
    - large_string: 12.20 µs -> 8.66 µs (-29.0%)

  scalar args, long strings (size=4096, count=64, strlen=128):
    - string_view:  40.13 µs -> 41.89 µs (+4.4%)
    - string:       46.96 µs -> 32.44 µs (-30.9%)
    - large_string: 47.24 µs -> 32.49 µs (-31.2%)

This PR doesn't modify the `string_view` code path; I've included the
benchmark results above for completeness, but any changes should just be
benchmarking noise.

## What changes are included in this PR?

* Implement optimization
* Other minor code cleanup
* Add a benchmark (only somewhat related to this optimization but
related to future optimization work)

## Are these changes tested?

Yes.

## Are there any user-facing changes?

No.

Author: neilconway

datafusion/functions/benches/substr.rs

@@ -201,7 +201,7 @@ fn criterion_benchmark(c: &mut Criterion) {
 
         group.finish();
 
-        // Scalar start, no count, long strings
+        // Scalar start, no count, long strings, start near middle
         let len = 128;
         let mut group = c.benchmark_group("substr, scalar start, no count, long strings");
         group.sampling_mode(SamplingMode::Flat);
@@ -220,6 +220,26 @@ fn criterion_benchmark(c: &mut Criterion) {
 
         group.finish();
 
+        // Scalar start=1, no count, long strings
+        let len = 128;
+        let mut group =
+            c.benchmark_group("substr, scalar start=1, no count, long strings");
+        group.sampling_mode(SamplingMode::Flat);
+        group.sample_size(10);
+
+        let args = create_args_without_count::<i32>(size, len, false, true, true);
+        group.bench_function(
+            format!("substr_string_view [size={size}, strlen={len}]"),
+            |b| b.iter(|| black_box(invoke_substr_with_args(args.clone(), size))),
+        );
+
+        let args = create_args_without_count::<i32>(size, len, false, false, true);
+        group.bench_function(format!("substr_string [size={size}, strlen={len}]"), |b| {
+            b.iter(|| black_box(invoke_substr_with_args(args.clone(), size)))
+        });
+
+        group.finish();
+
         // Scalar start and count, short strings
         let len = 12;
         let count = 6;

datafusion/functions/src/strings.rs

@@ -424,7 +424,7 @@ impl LargeStringArrayBuilder {
 /// - start_offset: The start offset of the substring in the view
 ///
 /// LLVM is apparently overly eager to inline this function into some hot loops,
-/// which bloats them and regresses performance, so we disable inling for now.
+/// which bloats them and regresses performance, so we disable inlining for now.
 #[inline(never)]
 pub fn append_view(
     views_buffer: &mut Vec<u128>,

datafusion/functions/src/unicode/substr.rs

@@ -20,8 +20,8 @@ use std::sync::Arc;
 use crate::strings::append_view;
 use crate::utils::make_scalar_function;
 use arrow::array::{
-    Array, ArrayRef, AsArray, Int64Array, StringArrayType, StringViewArray,
-    StringViewBuilder,
+    Array, ArrayRef, AsArray, GenericStringArray, Int64Array, OffsetSizeTrait,
+    StringArrayType, StringViewArray, StringViewBuilder, make_view,
 };
 use arrow::buffer::{NullBuffer, ScalarBuffer};
 use arrow::datatypes::DataType;
@@ -134,11 +134,11 @@ fn substr(args: &[ArrayRef]) -> Result<ArrayRef> {
     match args[0].data_type() {
         DataType::Utf8 => {
             let string_array = args[0].as_string::<i32>();
-            string_substr::<_>(string_array, &args[1..])
+            generic_string_substr(string_array, &args[1..])
         }
         DataType::LargeUtf8 => {
             let string_array = args[0].as_string::<i64>();
-            string_substr::<_>(string_array, &args[1..])
+            generic_string_substr(string_array, &args[1..])
         }
         DataType::Utf8View => {
             let string_array = args[0].as_string_view();
@@ -275,7 +275,7 @@ fn string_view_substr(
     let start_array = as_int64_array(&args[0])?;
     let count_array_opt = args.get(1).map(|a| as_int64_array(a)).transpose()?;
 
-    let enable_ascii_fast_path =
+    let is_ascii =
         enable_ascii_fast_path(&string_view_array, start_array, count_array_opt);
 
     // Combine null bitmaps from all inputs in bulk.
@@ -296,11 +296,10 @@ fn string_view_substr(
         let start = start_array.value(i);
         let count = count_array_opt.map(|a| a.value(i));
 
-        let (start, end) =
-            get_true_start_end(string, start, count, enable_ascii_fast_path)?;
-        let substr = &string[start..end];
+        let (byte_start, byte_end) = get_true_start_end(string, start, count, is_ascii)?;
+        let substr = &string[byte_start..byte_end];
 
-        append_view(&mut views_buf, raw_view, substr, start as u32);
+        append_view(&mut views_buf, raw_view, substr, byte_start as u32);
     }
 
     let views_buf = ScalarBuffer::from(views_buf);
@@ -319,15 +318,104 @@ fn string_view_substr(
     }
 }
 
-fn string_substr<'a, V>(string_array: V, args: &[ArrayRef]) -> Result<ArrayRef>
-where
-    V: StringArrayType<'a> + Copy,
-{
+fn values_fit_in_i32<T: OffsetSizeTrait>(string_array: &GenericStringArray<T>) -> bool {
+    // The Arrow spec defines StringView offset fields as signed 32-bit
+    // integers, so the maximum representable offset is i32::MAX.
+    string_array
+        .offsets()
+        .last()
+        .map(|offset| offset.as_usize() <= i32::MAX as usize)
+        .unwrap_or(true)
+}
+
+#[inline]
+fn append_view_from_buffer(
+    views_buf: &mut Vec<u128>,
+    substr: &str,
+    byte_offset: usize,
+) -> bool {
+    let byte_offset =
+        u32::try_from(byte_offset).expect("validated string buffer offset fits in i32");
+    let view = make_view(substr.as_bytes(), 0, byte_offset);
+    views_buf.push(view);
+    substr.len() > 12
+}
+
+#[expect(clippy::needless_range_loop)]
+fn generic_string_substr<T: OffsetSizeTrait>(
+    string_array: &GenericStringArray<T>,
+    args: &[ArrayRef],
+) -> Result<ArrayRef> {
+    // We'd like to return a StringViewArray that points into the input string
+    // array's values buffer. Since the Arrow spec defines StringView offsets
+    // as i32, we can't use this approach when the values buffer is >2GB, so
+    // fallback to copying.
+    if !values_fit_in_i32(string_array) {
+        return generic_string_substr_copy(string_array, args);
+    }
+
     let start_array = as_int64_array(&args[0])?;
     let count_array_opt = args.get(1).map(|a| as_int64_array(a)).transpose()?;
 
-    let enable_ascii_fast_path =
-        enable_ascii_fast_path(&string_array, start_array, count_array_opt);
+    let is_ascii = enable_ascii_fast_path(&string_array, start_array, count_array_opt);
+    let offsets = string_array.value_offsets();
+    let mut views_buf = Vec::with_capacity(string_array.len());
+    let mut has_out_of_line = false;
+
+    // Combine null bitmaps from all inputs in bulk.
+    let nulls = NullBuffer::union(
+        NullBuffer::union(string_array.nulls(), start_array.nulls()).as_ref(),
+        count_array_opt.and_then(|a| a.nulls()),
+    );
+
+    for i in 0..string_array.len() {
+        if nulls.as_ref().is_some_and(|n| n.is_null(i)) {
+            views_buf.push(0);
+            continue;
+        }
+
+        let string = string_array.value(i);
+        let source_offset = offsets[i].as_usize();
+        let start = start_array.value(i);
+        let count = count_array_opt.map(|a| a.value(i));
+
+        let (byte_start, byte_end) = get_true_start_end(string, start, count, is_ascii)?;
+        has_out_of_line |= append_view_from_buffer(
+            &mut views_buf,
+            &string[byte_start..byte_end],
+            source_offset + byte_start,
+        );
+    }
+
+    let views_buf = ScalarBuffer::from(views_buf);
+
+    // If all result strings are stored inline, we don't need to retain the
+    // input string array.
+    let data_buffers = if has_out_of_line {
+        vec![string_array.values().clone()]
+    } else {
+        vec![]
+    };
+
+    // Safety:
+    // (1) The blocks of the given views are all provided
+    // (2) Each referenced range in the source values buffer is within bounds
+    unsafe {
+        let array = StringViewArray::new_unchecked(views_buf, data_buffers, nulls);
+        Ok(Arc::new(array) as ArrayRef)
+    }
+}
+
+// Fallback for `generic_string_substr` if we can't use zerocopy because the
+// input string array is too large.
+fn generic_string_substr_copy<T: OffsetSizeTrait>(
+    string_array: &GenericStringArray<T>,
+    args: &[ArrayRef],
+) -> Result<ArrayRef> {
+    let start_array = as_int64_array(&args[0])?;
+    let count_array_opt = args.get(1).map(|a| as_int64_array(a)).transpose()?;
+
+    let is_ascii = enable_ascii_fast_path(&string_array, start_array, count_array_opt);
 
     // Combine null bitmaps from all inputs in bulk.
     let nulls = NullBuffer::union(
@@ -347,17 +435,20 @@ where
         let start = start_array.value(i);
         let count = count_array_opt.map(|a| a.value(i));
 
-        let (start, end) =
-            get_true_start_end(string, start, count, enable_ascii_fast_path)?;
-        result_builder.append_value(&string[start..end]);
+        let (byte_start, byte_end) = get_true_start_end(string, start, count, is_ascii)?;
+        result_builder.append_value(&string[byte_start..byte_end]);
     }
 
     Ok(Arc::new(result_builder.finish()) as ArrayRef)
 }
 
 #[cfg(test)]
 mod tests {
-    use arrow::array::{Array, StringViewArray};
+    use std::sync::Arc;
+
+    use arrow::array::{
+        Array, ArrayRef, AsArray, Int64Array, StringArray, StringViewArray,
+    };
     use arrow::datatypes::DataType::Utf8View;
 
     use datafusion_common::{Result, ScalarValue, exec_err};
@@ -734,4 +825,25 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn test_sliced_string_array_array_args() -> Result<()> {
+        // Use strings longer than 12 bytes so the result views are out-of-line.
+        let string_array = Arc::new(StringArray::from(vec![
+            "skipped_prefix_value",
+            "alphabet_long_string",
+            "joséésojanother_long",
+        ])) as ArrayRef;
+        let string_array = string_array.slice(1, 2);
+        let start_array = Arc::new(Int64Array::from(vec![3, 5])) as ArrayRef;
+        let count_array = Arc::new(Int64Array::from(vec![15, 14])) as ArrayRef;
+
+        let result = super::substr(&[string_array, start_array, count_array])?;
+        let result = result.as_string_view();
+
+        assert_eq!(result.value(0), "phabet_long_str");
+        assert_eq!(result.value(1), "ésojanother_lo");
+
+        Ok(())
+    }
 }