From 3e80e814c49af02137bf6c032b1145b35646fdc9 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 25 Jul 2026 07:52:23 -0600 Subject: [PATCH 1/2] 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) }); } From 4a7577fe70d7deb3ec407a24cb0b3a7939f1dc1f Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 27 Jul 2026 14:16:15 -0600 Subject: [PATCH 2/2] perf: size quote output buffer from the sliced array length `value_data().len()` spans the entire values buffer even when the array is a slice, and `get_buffer_memory_size()` reports buffer capacities while ignoring inlined view values. Take the length from the sliced offsets and from `total_bytes_len()` instead. --- datafusion/spark/src/function/string/quote.rs | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/datafusion/spark/src/function/string/quote.rs b/datafusion/spark/src/function/string/quote.rs index 548e2a2a4b8ef..38db82adfac90 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::{Array, ArrayRef, OffsetSizeTrait, StringBuilder}; +use arrow::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}; @@ -89,15 +89,22 @@ fn spark_quote_inner(arg: &[ArrayRef]) -> Result { fn quote_array(array: &ArrayRef) -> Result { let str_array = as_generic_string_array::(array)?; - Ok(quote_impl(str_array.iter(), str_array.value_data().len())) + // 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()`. + let offsets = str_array.value_offsets(); + let data_len = match (offsets.first(), offsets.last()) { + (Some(first), Some(last)) => last.as_usize() - first.as_usize(), + _ => 0, + }; + Ok(quote_impl(str_array.iter(), data_len)) } fn quote_view(str_view: &ArrayRef) -> Result { let str_array = as_string_view_array(str_view)?; - Ok(quote_impl( - str_array.iter(), - str_array.get_buffer_memory_size(), - )) + // `total_bytes_len` walks the (sliced) views and counts inlined values, + // unlike the buffer capacities reported by `get_buffer_memory_size`. + Ok(quote_impl(str_array.iter(), str_array.total_bytes_len())) } const QUOTE_CHAR: char = '\'';