diff --git a/datafusion/functions/src/string/common.rs b/datafusion/functions/src/string/common.rs index 7a5d2573ec257..b46353b609f1a 100644 --- a/datafusion/functions/src/string/common.rs +++ b/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( } } +/// 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( + string_array: &GenericStringArray, + nulls: Option, + 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 = Vec::with_capacity(end - start); + let mut offsets: Vec = 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::::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(args: &[ArrayRef]) -> Result { // Trim spaces by default - let result = string_array - .iter() - .map(|string| string.map(|s| Tr::trim_ascii_char(s, b' ').0)) - .collect::>(); - - 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::(&args[1])?; @@ -293,29 +351,31 @@ fn string_trim(args: &[ArrayRef]) -> Result = characters_array.value(0).chars().collect(); - let result = string_array - .iter() - .map(|item| item.map(|s| Tr::trim(s, &pattern).0)) - .collect::>(); - 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 = 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::>(); - - 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!(