perf: optimize string_trim

datafusion/functions/src/string/common.rs

@@ -27,7 +27,7 @@ use arrow::array::{
     Array, ArrayRef, AsArray, GenericStringArray, NullBufferBuilder, OffsetSizeTrait,
     StringViewArray, new_null_array,
 };
-use arrow::buffer::{Buffer, ScalarBuffer};
+use arrow::buffer::{Buffer, NullBuffer, OffsetBuffer, ScalarBuffer};
 use arrow::datatypes::DataType;
 use datafusion_common::Result;
 use datafusion_common::cast::{as_generic_string_array, as_string_view_array};
@@ -262,6 +262,65 @@ fn trim_and_append_view<Tr: Trimmer>(
     }
 }
 
+/// Builds the trimmed output array by writing the trimmed slices straight into
+/// the value buffer, rather than collecting through a string builder.
+///
+/// Every trimmed value is a substring of its input, so the byte range the input
+/// spans bounds the output's. Reserving that much up front means one allocation
+/// and no growth during the copy, and it also guarantees the running offset stays
+/// within `T` (the input array's own offsets already fit).
+///
+/// `nulls` becomes the output null buffer; null rows contribute no bytes.
+/// `trim_row` is called only for non-null rows, with the row index and its value,
+/// and must return a subslice of the value it is given.
+fn build_trimmed<T: OffsetSizeTrait, F>(
+    string_array: &GenericStringArray<T>,
+    nulls: Option<NullBuffer>,
+    mut trim_row: F,
+) -> ArrayRef
+where
+    F: for<'a> FnMut(usize, &'a str) -> &'a str,
+{
+    let len = string_array.len();
+    let input_offsets = string_array.value_offsets();
+    let start = input_offsets.first().unwrap().as_usize();
+    let end = input_offsets.last().unwrap().as_usize();
+
+    let mut values: Vec<u8> = Vec::with_capacity(end - start);
+    let mut offsets: Vec<T> = Vec::with_capacity(len + 1);
+    offsets.push(T::usize_as(0));
+
+    match &nulls {
+        // Keeping the null check out of the all-valid path leaves it branch-free.
+        None => {
+            for i in 0..len {
+                // SAFETY: `i` is in bounds.
+                let s = unsafe { string_array.value_unchecked(i) };
+                values.extend_from_slice(trim_row(i, s).as_bytes());
+                offsets.push(T::usize_as(values.len()));
+            }
+        }
+        Some(validity) => {
+            for i in 0..len {
+                if validity.is_valid(i) {
+                    // SAFETY: `i` is in bounds.
+                    let s = unsafe { string_array.value_unchecked(i) };
+                    values.extend_from_slice(trim_row(i, s).as_bytes());
+                }
+                offsets.push(T::usize_as(values.len()));
+            }
+        }
+    }
+
+    let offsets = OffsetBuffer::new(ScalarBuffer::from(offsets));
+    // SAFETY: trimming splits `s` on char boundaries, so the value buffer is a
+    // concatenation of valid UTF-8; the offsets are monotonic and end at its length.
+    let array = unsafe {
+        GenericStringArray::<T>::new_unchecked(offsets, Buffer::from_vec(values), nulls)
+    };
+    Arc::new(array)
+}
+
 /// Applies the trim function to the given string array(s)
 /// and returns a new string array with the trimmed values.
 ///
@@ -273,12 +332,11 @@ fn string_trim<T: OffsetSizeTrait, Tr: Trimmer>(args: &[ArrayRef]) -> Result<Arr
     match args.len() {
         1 => {
             // Trim spaces by default
-            let result = string_array
-                .iter()
-                .map(|string| string.map(|s| Tr::trim_ascii_char(s, b' ').0))
-                .collect::<GenericStringArray<T>>();
-
-            Ok(Arc::new(result) as ArrayRef)
+            Ok(build_trimmed(
+                string_array,
+                string_array.nulls().cloned(),
+                |_, s| Tr::trim_ascii_char(s, b' ').0,
+            ))
         }
         2 => {
             let characters_array = as_generic_string_array::<T>(&args[1])?;
@@ -293,29 +351,31 @@ fn string_trim<T: OffsetSizeTrait, Tr: Trimmer>(args: &[ArrayRef]) -> Result<Arr
                 }
 
                 let pattern: Vec<char> = characters_array.value(0).chars().collect();
-                let result = string_array
-                    .iter()
-                    .map(|item| item.map(|s| Tr::trim(s, &pattern).0))
-                    .collect::<GenericStringArray<T>>();
-                return Ok(Arc::new(result) as ArrayRef);
+                return Ok(build_trimmed(
+                    string_array,
+                    string_array.nulls().cloned(),
+                    |_, s| Tr::trim(s, &pattern).0,
+                ));
+            }
+
+            // Indexing `characters_array` per row below requires the two arguments
+            // to line up.
+            if characters_array.len() != string_array.len() {
+                return exec_err!(
+                    "Function TRIM was called with mismatched argument lengths"
+                );
             }
 
+            // A row is null if either argument is null.
+            let nulls = NullBuffer::union(string_array.nulls(), characters_array.nulls());
+
             // Per-row pattern - must compute pattern chars for each row
             let mut pattern: Vec<char> = Vec::new();
-            let result = string_array
-                .iter()
-                .zip(characters_array.iter())
-                .map(|(string, characters)| match (string, characters) {
-                    (Some(s), Some(c)) => {
-                        pattern.clear();
-                        pattern.extend(c.chars());
-                        Some(Tr::trim(s, &pattern).0)
-                    }
-                    _ => None,
-                })
-                .collect::<GenericStringArray<T>>();
-
-            Ok(Arc::new(result) as ArrayRef)
+            Ok(build_trimmed(string_array, nulls, |i, s| {
+                pattern.clear();
+                pattern.extend(characters_array.value(i).chars());
+                Tr::trim(s, &pattern).0
+            }))
         }
         other => {
             exec_err!(