perf: optimize spark soundex function

datafusion/spark/Cargo.toml

@@ -105,3 +105,7 @@ name = "sha2"
 [[bench]]
 harness = false
 name = "floor"
+
+[[bench]]
+harness = false
+name = "soundex"

datafusion/spark/benches/soundex.rs

@@ -0,0 +1,71 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow::array::{ArrayRef, StringArray};
+use criterion::{Criterion, criterion_group, criterion_main};
+use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl};
+use datafusion_spark::function::string::soundex::SparkSoundex;
+use std::hint::black_box;
+use std::sync::Arc;
+
+fn create_string_array(rows: usize) -> ArrayRef {
+    let strings: Vec<String> = (0..rows)
+        .map(|i| match i % 5 {
+            0 => format!("Robert{i}"),
+            1 => format!("Rupert{i}"),
+            2 => format!("Washington{i}"),
+            3 => format!("123Invalid{i}"),
+            _ => format!("Short{i}"),
+        })
+        .collect();
+
+    Arc::new(StringArray::from(
+        strings.iter().map(|s| s.as_str()).collect::<Vec<&str>>(),
+    )) as ArrayRef
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+    let rows = 8192;
+    let array = create_string_array(rows);
+    let soundex_udf = SparkSoundex::new();
+
+    c.bench_function("spark_soundex: standard utf8 array", |b| {
+        let args = vec![ColumnarValue::Array(Arc::clone(&array))];
+        b.iter(|| {
+            black_box(
+                soundex_udf
+                    .invoke_with_args(ScalarFunctionArgs {
+                        args: args.clone(),
+                        arg_fields: vec![],
+                        number_rows: rows,
+                        return_field: Arc::new(arrow::datatypes::Field::new(
+                            "soundex",
+                            arrow::datatypes::DataType::Utf8,
+                            true,
+                        )),
+                        config_options: Arc::new(
+                            datafusion::config::ConfigOptions::default(),
+                        ),
+                    })
+                    .unwrap(),
+            )
+        })
+    });
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);

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

@@ -15,7 +15,8 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use arrow::array::{ArrayRef, OffsetSizeTrait, StringArray};
+use arrow::array::StringBuilder;
+use arrow::array::{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;
@@ -79,22 +80,54 @@ 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))
+fn soundex_array<O: OffsetSizeTrait>(array: &ArrayRef) -> Result<ArrayRef> {
+    let str_array = as_generic_string_array::<O>(array)?;
+
+    // Pre-allocate exact memory: row count and total string bytes (each soundex is 4 bytes)
+    let mut builder =
+        GenericStringBuilder::<O>::with_capacity(str_array.len(), str_array.len() * 4);
+
+    for opt_s in str_array.iter() {
+        if let Some(s) = opt_s {
+            let code = compute_soundex(s);
+            builder.append_value(&code);
+        } else {
+            builder.append_null();
+        }
+    }
+
+    Ok(Arc::new(builder.finish()) as ArrayRef)
 }
 
 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)
+
+    // Pre-allocate for Utf8View as well
+    let mut builder = StringBuilder::with_capacity(str_array.len(), str_array.len() * 4);
+
+    for opt_str in str_array.iter() {
+        if let Some(s) = opt_str {
+            let code = compute_soundex(s);
+            builder.append_value(&code);
+        } else {
+            builder.append_null();
+        }
+    }
+
+    Ok(Arc::new(builder.finish()) as ArrayRef)
+}
+
+#[inline]
+fn classify_byte(c: u8) -> Option<u8> {
+    match c {
+        b'B' | b'F' | b'P' | b'V' => Some(b'1'),
+        b'C' | b'G' | b'J' | b'K' | b'Q' | b'S' | b'X' | b'Z' => Some(b'2'),
+        b'D' | b'T' => Some(b'3'),
+        b'L' => Some(b'4'),
+        b'M' | b'N' => Some(b'5'),
+        b'R' => Some(b'6'),
+        _ => None,
+    }
 }
 
 fn classify_char(c: char) -> Option<char> {
@@ -114,30 +147,74 @@ fn is_ignored(c: char) -> bool {
 }
 
 fn compute_soundex(s: &str) -> String {
-    let mut chars = s.chars();
+    // Fast path for ASCII strings: operate directly on raw bytes with zero heap allocations
+    if s.is_ascii() {
+        let bytes = s.as_bytes();
+        let mut iter = bytes.iter();
+
+        let first_byte = match iter.next() {
+            Some(&c) if c.is_ascii_alphabetic() => c.to_ascii_uppercase(),
+            _ => return s.to_string(),
+        };
+
+        // Soundex codes are always exactly 4 characters, padded with '0'
+        let mut buf = [b'0'; 4];
+        buf[0] = first_byte;
+        let mut len = 1;
+        let mut last_code = classify_byte(first_byte);
+
+        for &c in iter {
+            if len >= 4 {
+                break;
+            }
+            let upper = c.to_ascii_uppercase();
+            if upper == b'H' || upper == b'W' {
+                continue;
+            }
 
+            match classify_byte(upper) {
+                Some(code) => {
+                    if last_code != Some(code) {
+                        buf[len] = code;
+                        len += 1;
+                    }
+                    last_code = Some(code);
+                }
+                None => {
+                    last_code = None;
+                }
+            }
+        }
+
+        // SAFETY: buf contains exclusively valid ASCII uppercase letters and '0' digits
+        return unsafe { std::str::from_utf8_unchecked(&buf).to_string() };
+    }
+
+    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(),
     };
 
-    let mut soundex_code = String::with_capacity(4);
-    soundex_code.push(first_char);
+    let mut buf = [b'0'; 4];
+    buf[0] = first_char as u8;
+    let mut len = 1;
     let mut last_code = classify_char(first_char);
 
     for c in chars {
-        if soundex_code.len() >= 4 {
+        if len >= 4 {
             break;
         }
-
         if is_ignored(c) {
             continue;
         }
 
         match classify_char(c) {
             Some(code) => {
+                let code_u8 = code as u8;
                 if last_code != Some(code) {
-                    soundex_code.push(code);
+                    buf[len] = code_u8;
+                    len += 1;
                 }
                 last_code = Some(code);
             }
@@ -146,5 +223,6 @@ fn compute_soundex(s: &str) -> String {
             }
         }
     }
-    format!("{soundex_code:0<4}")
+
+    unsafe { std::str::from_utf8_unchecked(&buf).to_string() }
 }