From 3e80e814c49af02137bf6c032b1145b35646fdc9 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 25 Jul 2026 07:52:23 -0600 Subject: [PATCH] perf: remove per-row String allocation from Spark soundex and quote Both functions built a fresh String for every row and collected the results into a StringArray. soundex allocated twice per row -- once for the code buffer and once more for the format! that zero-pads it -- and quote allocated a String sized to the input before copying it in character at a time. Neither needs to allocate. A soundex code is always exactly four ASCII characters, so it is built in a stack buffer. quote writes straight into the builder and copies the runs between quotes rather than one char at a time. soundex -50%, quote -61% against the benchmarks added in #23882. --- datafusion/spark/src/function/string/quote.rs | 64 +++++++++++++------ .../spark/src/function/string/soundex.rs | 58 +++++++++++------ 2 files changed, 81 insertions(+), 41 deletions(-) diff --git a/datafusion/spark/src/function/string/quote.rs b/datafusion/spark/src/function/string/quote.rs index 39ad8bf841764..548e2a2a4b8ef 100644 --- a/datafusion/spark/src/function/string/quote.rs +++ b/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::{Array, ArrayRef, OffsetSizeTrait, StringBuilder}; use arrow::datatypes::DataType; use datafusion::logical_expr::{Coercion, ColumnarValue, Signature, TypeSignatureClass}; use datafusion_common::cast::{as_generic_string_array, as_string_view_array}; @@ -25,6 +25,7 @@ use datafusion_common::{Result, exec_err}; use datafusion_expr::{ScalarFunctionArgs, ScalarUDFImpl, Volatility}; use datafusion_functions::utils::make_scalar_function; +use std::fmt::Write; use std::sync::Arc; /// Spark-compatible `quote` expression @@ -88,34 +89,55 @@ fn spark_quote_inner(arg: &[ArrayRef]) -> Result { fn quote_array(array: &ArrayRef) -> Result { let str_array = as_generic_string_array::(array)?; - let result = str_array - .iter() - .map(|s| s.map(compute_quote)) - .collect::(); - Ok(Arc::new(result)) + Ok(quote_impl(str_array.iter(), str_array.value_data().len())) } fn quote_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_quote)) - .collect::(); - Ok(Arc::new(result) as ArrayRef) + Ok(quote_impl( + str_array.iter(), + str_array.get_buffer_memory_size(), + )) } const QUOTE_CHAR: char = '\''; -const ESCAPE_CHAR: char = '\\'; +/// A literal quote in the input is emitted as this two-character escape. +const ESCAPED_QUOTE: &str = "\\'"; -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); +/// 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>( + input: impl Iterator>, + data_capacity: usize, +) -> ArrayRef { + let len = input.size_hint().0; + let mut builder = StringBuilder::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(builder: &mut StringBuilder, 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(""); } diff --git a/datafusion/spark/src/function/string/soundex.rs b/datafusion/spark/src/function/string/soundex.rs index 1fef0d5384821..c2878fbf22809 100644 --- a/datafusion/spark/src/function/string/soundex.rs +++ b/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, OffsetSizeTrait, StringBuilder}; 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,24 @@ 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)) + Ok(soundex_impl(as_generic_string_array::(array)?.iter())) } 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) + Ok(soundex_impl(as_string_view_array(str_view)?.iter())) +} + +fn soundex_impl<'a>(input: impl Iterator>) -> ArrayRef { + let len = input.size_hint().0; + // A soundex code is always exactly 4 ASCII characters. + let mut builder = StringBuilder::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 { @@ -113,20 +116,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(builder: &mut StringBuilder, 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 +152,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 +162,7 @@ fn compute_soundex(s: &str) -> String { } } } - format!("{soundex_code:0<4}") + + // SAFETY: `soundex_code` holds an ASCII letter followed by ASCII digits. + builder.append_value(unsafe { std::str::from_utf8_unchecked(&soundex_code) }); }