Function string_trim

datafusion/functions/src/string/common.rs:270–326

Applies the trim function to the given string array(s) and returns a new string array with the trimmed values. Pre-computes the pattern characters once for scalar patterns to avoid repeated allocations per row.

(args: &[ArrayRef])

/// Pre-computes the pattern characters once for scalar patterns to avoid
/// repeated allocations per row.
fn string_trim<T: OffsetSizeTrait, Tr: Trimmer>(args: &[ArrayRef]) -> Result<ArrayRef> {
    let string_array = as_generic_string_array::<T>(&args[0])?;

    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)
        }
        2 => {
            let characters_array = as_generic_string_array::<T>(&args[1])?;

            if characters_array.len() == 1 {
                // Scalar pattern - pre-compute pattern chars once
                if characters_array.is_null(0) {
                    return Ok(new_null_array(
                        string_array.data_type(),
                        string_array.len(),
                    ));
                }

                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);
            }

            // 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)
        }
        other => {
            exec_err!(
                "Function TRIM was called with {other} arguments. It requires at least 1 and at most 2."
            )
        }
    }
}