perf: remove per-row String allocation from Spark soundex and quote

datafusion/spark/src/function/string/quote.rs

@@ -15,7 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use arrow::array::{ArrayRef, OffsetSizeTrait, StringArray};
+use arrow::array::{ArrayRef, GenericStringBuilder, OffsetSizeTrait};
 use arrow::datatypes::DataType;
 use datafusion_common::cast::{as_generic_string_array, as_string_view_array};
 use datafusion_common::types::{NativeType, logical_string};
@@ -27,6 +27,7 @@ use datafusion_expr::{
 };
 use datafusion_functions::utils::make_scalar_function;
 
+use std::fmt::Write;
 use std::sync::Arc;
 
 /// Spark-compatible `quote` expression
@@ -90,34 +91,139 @@ fn spark_quote_inner(arg: &[ArrayRef]) -> Result<ArrayRef> {
 
 fn quote_array<T: OffsetSizeTrait>(array: &ArrayRef) -> Result<ArrayRef> {
     let str_array = as_generic_string_array::<T>(array)?;
-    let result = str_array
-        .iter()
-        .map(|s| s.map(compute_quote))
-        .collect::<StringArray>();
-    Ok(Arc::new(result))
+    // Slicing an array keeps the whole value buffer and narrows only the
+    // offsets, so measure the data through the offsets rather than through
+    // `value_data()`. An offset buffer holds one more entry than the array has
+    // rows, so it is never empty.
+    let offsets = str_array.value_offsets();
+    let data_len = offsets.last().unwrap().as_usize() - offsets[0].as_usize();
+    Ok(quote_impl::<T, _>(str_array.iter(), data_len))
 }
 
 fn quote_view(str_view: &ArrayRef) -> Result<ArrayRef> {
     let str_array = as_string_view_array(str_view)?;
-    let result = str_array
-        .iter()
-        .map(|opt_str| opt_str.map(compute_quote))
-        .collect::<StringArray>();
-    Ok(Arc::new(result) as ArrayRef)
+    // `total_bytes_len` walks the (sliced) views and counts inlined values,
+    // unlike the buffer capacities reported by `get_buffer_memory_size`.
+    Ok(quote_impl::<i32, _>(
+        str_array.iter(),
+        str_array.total_bytes_len(),
+    ))
 }
 
 const QUOTE_CHAR: char = '\'';
-const ESCAPE_CHAR: char = '\\';
-
-fn compute_quote(s: &str) -> String {
-    let mut quoted = String::with_capacity(s.len() + 2);
-    quoted.push(QUOTE_CHAR);
-    for c in s.chars() {
-        if c == QUOTE_CHAR {
-            quoted.push(ESCAPE_CHAR);
+/// A literal quote in the input is emitted as this two-character escape.
+const ESCAPED_QUOTE: &str = "\\'";
+
+/// Quotes every value, writing directly into the output buffer.
+///
+/// `data_capacity` is a hint for the total input byte length; the output adds two
+/// surrounding quotes per row plus one byte per escaped quote.
+fn quote_impl<'a, O: OffsetSizeTrait, I: Iterator<Item = Option<&'a str>>>(
+    input: I,
+    data_capacity: usize,
+) -> ArrayRef {
+    let len = input.size_hint().0;
+    let mut builder =
+        GenericStringBuilder::<O>::with_capacity(len, data_capacity + 2 * len);
+    for value in input {
+        match value {
+            Some(value) => append_quoted(&mut builder, value),
+            None => builder.append_null(),
         }
-        quoted.push(c);
     }
-    quoted.push(QUOTE_CHAR);
-    quoted
+    Arc::new(builder.finish())
+}
+
+/// Appends `s` wrapped in single quotes, with any embedded quote backslash-escaped.
+///
+/// Writes straight into the builder's buffer — finalized by the trailing
+/// `append_value("")` — so no intermediate `String` is allocated per row, and
+/// copies the runs between quotes rather than one character at a time.
+fn append_quoted<O: OffsetSizeTrait>(builder: &mut GenericStringBuilder<O>, s: &str) {
+    // `write_str` on a `GenericStringBuilder` is infallible.
+    let mut runs = s.split(QUOTE_CHAR);
+    builder.write_char(QUOTE_CHAR).unwrap();
+    // `split` always yields at least one run.
+    builder.write_str(runs.next().unwrap_or_default()).unwrap();
+    for run in runs {
+        builder.write_str(ESCAPED_QUOTE).unwrap();
+        builder.write_str(run).unwrap();
+    }
+    builder.write_char(QUOTE_CHAR).unwrap();
+    builder.append_value("");
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use arrow::array::{LargeStringArray, StringArray, StringViewArray};
+
+    fn quote(array: ArrayRef) -> ArrayRef {
+        spark_quote_inner(&[array]).unwrap()
+    }
+
+    fn as_strings(array: &ArrayRef) -> Vec<Option<&str>> {
+        as_generic_string_array::<i32>(array)
+            .unwrap()
+            .iter()
+            .collect()
+    }
+
+    #[test]
+    fn quote_preserves_the_offset_width() {
+        let utf8 = quote(Arc::new(StringArray::from(vec!["it's"])) as ArrayRef);
+        assert_eq!(utf8.data_type(), &DataType::Utf8);
+        assert_eq!(as_strings(&utf8), vec![Some("'it\\'s'")]);
+
+        let large = quote(Arc::new(LargeStringArray::from(vec!["it's"])) as ArrayRef);
+        assert_eq!(large.data_type(), &DataType::LargeUtf8);
+        let large = as_generic_string_array::<i64>(&large).unwrap();
+        assert_eq!(large.value(0), "'it\\'s'");
+
+        // A view input has no offsets to preserve, so it narrows to `Utf8`.
+        let view = quote(Arc::new(StringViewArray::from(vec!["it's"])) as ArrayRef);
+        assert_eq!(view.data_type(), &DataType::Utf8);
+        assert_eq!(as_strings(&view), vec![Some("'it\\'s'")]);
+    }
+
+    /// Slicing keeps the whole value buffer, so the capacity hint has to be
+    /// measured through the offsets rather than through `value_data()`.
+    #[test]
+    fn quote_sliced_array() {
+        let array = Arc::new(StringArray::from(vec![
+            Some("a very long leading value that inflates value_data"),
+            Some("it's"),
+            None,
+            Some(""),
+        ])) as ArrayRef;
+
+        let result = quote(array.slice(1, 3));
+        assert_eq!(
+            as_strings(&result),
+            vec![Some("'it\\'s'"), None, Some("''")]
+        );
+    }
+
+    #[test]
+    fn quote_sliced_view_array() {
+        let array = Arc::new(StringViewArray::from(vec![
+            // Longer than 12 bytes, so this one lives in a data buffer.
+            Some("a very long leading value that inflates value_data"),
+            Some("it's"),
+            None,
+        ])) as ArrayRef;
+
+        let result = quote(array.slice(1, 2));
+        assert_eq!(as_strings(&result), vec![Some("'it\\'s'"), None]);
+    }
+
+    #[test]
+    fn quote_empty_array() {
+        let empty = quote(Arc::new(StringArray::from(Vec::<&str>::new())) as ArrayRef);
+        assert_eq!(empty.len(), 0);
+
+        let empty_view =
+            quote(Arc::new(StringViewArray::from(Vec::<&str>::new())) as ArrayRef);
+        assert_eq!(empty_view.len(), 0);
+    }
 }

datafusion/spark/src/function/string/soundex.rs

@@ -15,7 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use arrow::array::{ArrayRef, OffsetSizeTrait, StringArray};
+use arrow::array::{ArrayRef, GenericStringBuilder, OffsetSizeTrait};
 use arrow::datatypes::DataType;
 use datafusion_common::cast::{as_generic_string_array, as_string_view_array};
 use datafusion_common::utils::take_function_args;
@@ -80,21 +80,30 @@ fn spark_soundex_inner(arg: &[ArrayRef]) -> Result<ArrayRef> {
 }
 
 fn soundex_array<T: OffsetSizeTrait>(array: &ArrayRef) -> Result<ArrayRef> {
-    let str_array = as_generic_string_array::<T>(array)?;
-    let result = str_array
-        .iter()
-        .map(|s| s.map(compute_soundex))
-        .collect::<StringArray>();
-    Ok(Arc::new(result))
+    Ok(soundex_impl::<T, _>(
+        as_generic_string_array::<T>(array)?.iter(),
+    ))
 }
 
 fn soundex_view(str_view: &ArrayRef) -> Result<ArrayRef> {
-    let str_array = as_string_view_array(str_view)?;
-    let result = str_array
-        .iter()
-        .map(|opt_str| opt_str.map(compute_soundex))
-        .collect::<StringArray>();
-    Ok(Arc::new(result) as ArrayRef)
+    Ok(soundex_impl::<i32, _>(
+        as_string_view_array(str_view)?.iter(),
+    ))
+}
+
+fn soundex_impl<'a, O: OffsetSizeTrait, I: Iterator<Item = Option<&'a str>>>(
+    input: I,
+) -> ArrayRef {
+    let len = input.size_hint().0;
+    // A soundex code is always exactly 4 ASCII characters.
+    let mut builder = GenericStringBuilder::<O>::with_capacity(len, len * SOUNDEX_LEN);
+    for value in input {
+        match value {
+            Some(value) => append_soundex(&mut builder, value),
+            None => builder.append_null(),
+        }
+    }
+    Arc::new(builder.finish())
 }
 
 fn classify_char(c: char) -> Option<char> {
@@ -113,20 +122,32 @@ fn is_ignored(c: char) -> bool {
     matches!(c.to_ascii_uppercase(), 'H' | 'W')
 }
 
-fn compute_soundex(s: &str) -> String {
+/// Length of a soundex code: an initial letter plus three digits.
+const SOUNDEX_LEN: usize = 4;
+
+/// Appends the soundex code of `s` to `builder`.
+///
+/// Strings that do not start with an ASCII letter are passed through unchanged.
+/// Otherwise the code is built in a stack buffer, so no row allocates.
+fn append_soundex<O: OffsetSizeTrait>(builder: &mut GenericStringBuilder<O>, s: &str) {
     let mut chars = s.chars();
 
     let first_char = match chars.next() {
         Some(c) if c.is_ascii_alphabetic() => c.to_ascii_uppercase(),
-        _ => return s.to_string(),
+        _ => {
+            builder.append_value(s);
+            return;
+        }
     };
 
-    let mut soundex_code = String::with_capacity(4);
-    soundex_code.push(first_char);
+    // Codes shorter than four characters are right-padded with '0'.
+    let mut soundex_code = [b'0'; SOUNDEX_LEN];
+    soundex_code[0] = first_char as u8;
+    let mut written = 1;
     let mut last_code = classify_char(first_char);
 
     for c in chars {
-        if soundex_code.len() >= 4 {
+        if written >= SOUNDEX_LEN {
             break;
         }
 
@@ -137,7 +158,8 @@ fn compute_soundex(s: &str) -> String {
         match classify_char(c) {
             Some(code) => {
                 if last_code != Some(code) {
-                    soundex_code.push(code);
+                    soundex_code[written] = code as u8;
+                    written += 1;
                 }
                 last_code = Some(code);
             }
@@ -146,5 +168,99 @@ fn compute_soundex(s: &str) -> String {
             }
         }
     }
-    format!("{soundex_code:0<4}")
+
+    // `soundex_code` holds an ASCII letter followed by ASCII digits, so the
+    // validation here is a four-byte check that never fails.
+    builder
+        .append_value(std::str::from_utf8(&soundex_code).expect("soundex code is ASCII"));
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use arrow::array::{LargeStringArray, StringArray, StringViewArray};
+
+    fn soundex(array: ArrayRef) -> ArrayRef {
+        spark_soundex_inner(&[array]).unwrap()
+    }
+
+    fn as_strings(array: &ArrayRef) -> Vec<Option<&str>> {
+        as_generic_string_array::<i32>(array)
+            .unwrap()
+            .iter()
+            .collect()
+    }
+
+    #[test]
+    fn soundex_preserves_the_offset_width() {
+        let utf8 = soundex(Arc::new(StringArray::from(vec!["Miller"])) as ArrayRef);
+        assert_eq!(utf8.data_type(), &DataType::Utf8);
+        assert_eq!(as_strings(&utf8), vec![Some("M460")]);
+
+        let large = soundex(Arc::new(LargeStringArray::from(vec!["Miller"])) as ArrayRef);
+        assert_eq!(large.data_type(), &DataType::LargeUtf8);
+        assert_eq!(
+            as_generic_string_array::<i64>(&large).unwrap().value(0),
+            "M460"
+        );
+
+        // A view input has no offsets to preserve, so it narrows to `Utf8`.
+        let view = soundex(Arc::new(StringViewArray::from(vec!["Miller"])) as ArrayRef);
+        assert_eq!(view.data_type(), &DataType::Utf8);
+        assert_eq!(as_strings(&view), vec![Some("M460")]);
+    }
+
+    /// Values whose first character is not an ASCII letter are passed through
+    /// unchanged, so the output is not always the four-byte code.
+    #[test]
+    fn soundex_multi_row_batch() {
+        let array = Arc::new(StringArray::from(vec![
+            Some("Miller"),
+            None,
+            Some(""),
+            // Non-ASCII alphabetic first character: passthrough, not a code.
+            Some("Ñoño"),
+            Some("Éclair"),
+            Some("123"),
+            Some("Robert"),
+        ])) as ArrayRef;
+
+        assert_eq!(
+            as_strings(&soundex(array)),
+            vec![
+                Some("M460"),
+                None,
+                Some(""),
+                Some("Ñoño"),
+                Some("Éclair"),
+                Some("123"),
+                Some("R163"),
+            ]
+        );
+    }
+
+    #[test]
+    fn soundex_sliced_array() {
+        let array = Arc::new(StringArray::from(vec![
+            Some("Miller"),
+            Some("Ñoño"),
+            None,
+            Some("Robert"),
+        ])) as ArrayRef;
+
+        assert_eq!(
+            as_strings(&soundex(array.slice(1, 3))),
+            vec![Some("Ñoño"), None, Some("R163")]
+        );
+    }
+
+    #[test]
+    fn soundex_empty_array() {
+        let empty = soundex(Arc::new(StringArray::from(Vec::<&str>::new())) as ArrayRef);
+        assert_eq!(empty.len(), 0);
+
+        let empty_view =
+            soundex(Arc::new(StringViewArray::from(Vec::<&str>::new())) as ArrayRef);
+        assert_eq!(empty_view.len(), 0);
+    }
 }