diff --git a/datafusion/spark/Cargo.toml b/datafusion/spark/Cargo.toml index 93987b553f2f5..ddbc7d92b4bc0 100644 --- a/datafusion/spark/Cargo.toml +++ b/datafusion/spark/Cargo.toml @@ -105,3 +105,7 @@ name = "sha2" [[bench]] harness = false name = "floor" + +[[bench]] +harness = false +name = "soundex" diff --git a/datafusion/spark/benches/soundex.rs b/datafusion/spark/benches/soundex.rs new file mode 100644 index 0000000000000..bc0af111eb2e2 --- /dev/null +++ b/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 = (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::>(), + )) 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); diff --git a/datafusion/spark/src/function/string/soundex.rs b/datafusion/spark/src/function/string/soundex.rs index 1fef0d5384821..391176f0ad21e 100644 --- a/datafusion/spark/src/function/string/soundex.rs +++ b/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 { } } -fn soundex_array(array: &ArrayRef) -> Result { - let str_array = as_generic_string_array::(array)?; - let result = str_array - .iter() - .map(|s| s.map(compute_soundex)) - .collect::(); - Ok(Arc::new(result)) +fn soundex_array(array: &ArrayRef) -> Result { + let str_array = as_generic_string_array::(array)?; + + // Pre-allocate exact memory: row count and total string bytes (each soundex is 4 bytes) + let mut builder = + GenericStringBuilder::::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 { let str_array = as_string_view_array(str_view)?; - let result = str_array - .iter() - .map(|opt_str| opt_str.map(compute_soundex)) - .collect::(); - 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 { + 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 { @@ -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() } }