You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/09/05 20:03:13 UTC
(datafusion) 01/01: perf: remove per-row String allocation from Spark soundex and quote (#23880)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-23880-33028d5543ed2fd2ad19bd247888cfb3c55cec3b
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit be5f2654ea68a38a7cf1ef78d20c70600727264a
Author: Andy Grove <ag...@apache.org>
AuthorDate: Sat Sep 5 20:02:51 2026 +0000
perf: remove per-row String allocation from Spark soundex and quote (#23880)
## Which issue does this PR close?
N/A
## Rationale for this change
Spark's `soundex` and `quote` both allocated a fresh `String` for every
row and
collected the results into a `StringArray`.
`soundex` allocated twice per row: once for the code buffer, and again
for the
`format!("{soundex_code:0<4}")` that zero-pads it. `quote` allocated a
`String`
sized to the input, then copied the input into it one `char` at a time.
Neither function needs to allocate per row. A soundex code is always
exactly
four ASCII characters, so it fits in a stack buffer. `quote` only ever
wraps the
input and escapes embedded quotes, so it can write straight into the
output
buffer and copy the runs between quotes rather than character by
character.
## What changes are included in this PR?
`soundex.rs`:
- `compute_soundex(&str) -> String` becomes `append_soundex(&mut
StringBuilder, &str)`,
building the four-character code in a `[u8; 4]` initialised to `b'0'` —
which is
the zero-padding, so the trailing `format!` disappears.
- Strings that do not start with an ASCII letter are passed through
unchanged;
previously this path called `s.to_string()`, now it appends by
reference.
- The `Utf8`/`LargeUtf8` and `Utf8View` entry points share one
`soundex_impl`
over `Option<&str>`, pre-sizing the builder at 4 bytes per row.
`quote.rs`:
- `compute_quote(&str) -> String` becomes `append_quoted(&mut
StringBuilder, &str)`,
writing into the builder's buffer via the `fmt::Write` impl and
finalising with
`append_value("")`.
- Escaping uses `str::split('\'')` to copy the runs between quotes in
one memcpy
each, instead of pushing every `char` individually.
- The builder is pre-sized from the input's `value_data()` length plus
two bytes
per row for the surrounding quotes.
Output is unchanged in both cases.
## Are these changes tested?
Existing coverage pins the behaviour: `spark/string/soundex.slt` and
`spark/string/quote.slt` assert concrete outputs, including the
non-alphabetic
passthrough, codes shorter than four characters (padding), codes
truncated at
four, and strings with and without embedded quotes. All 59
`spark/string`
sqllogictest files and the 258 `datafusion-spark` unit tests pass.
The benchmark used below, `datafusion/spark/benches/soundex_quote.rs`,
is added
separately in #23882 so the baseline can be measured on `main` before
this
change lands. It covers both functions over `Utf8` and `Utf8View` at
1024 and
8192 rows with 20% nulls; the `quote` input is a mix of strings with and
without
embedded quotes so the escaping path is exercised without dominating.
### Benchmarks
Criterion, `apache/main` @ `f1ab86dad` as baseline. Median of the
reported
change interval.
| Benchmark | 1024 | 8192 |
| --- | --- | --- |
| `soundex/utf8` | −52.9% | −49.5% |
| `soundex/utf8view` | −54.5% | −50.2% |
| `quote/utf8` | −60.7% | −61.5% |
| `quote/utf8view` | −60.5% | −61.3% |
## Are there any user-facing changes?
No. Both functions produce byte-identical output; this is purely an
allocation
change.
---
datafusion/spark/src/function/string/quote.rs | 150 +++++++++++++++++---
datafusion/spark/src/function/string/soundex.rs | 156 ++++++++++++++++++---
.../sqllogictest/test_files/spark/string/quote.slt | 61 ++++++++
.../test_files/spark/string/soundex.slt | 53 +++++++
4 files changed, 378 insertions(+), 42 deletions(-)
diff --git a/datafusion/spark/src/function/string/quote.rs b/datafusion/spark/src/function/string/quote.rs
index 55d827fb7a..b4fd5e3827 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::{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);
+ }
}
diff --git a/datafusion/spark/src/function/string/soundex.rs b/datafusion/spark/src/function/string/soundex.rs
index 1fef0d5384..8313a15b81 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, 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);
+ }
}
diff --git a/datafusion/sqllogictest/test_files/spark/string/quote.slt b/datafusion/sqllogictest/test_files/spark/string/quote.slt
index b5ef0f84e6..e63bf196d7 100644
--- a/datafusion/sqllogictest/test_files/spark/string/quote.slt
+++ b/datafusion/sqllogictest/test_files/spark/string/quote.slt
@@ -159,3 +159,64 @@ query T
SELECT quote('''');
----
'\''
+
+query T
+SELECT quote(NULL);
+----
+NULL
+
+query T
+SELECT quote('');
+----
+''
+
+# LargeUtf8 and Utf8View inputs
+
+query T
+SELECT quote(arrow_cast('it''s', 'LargeUtf8'));
+----
+'it\'s'
+
+query T
+SELECT quote(arrow_cast('', 'LargeUtf8'));
+----
+''
+
+query T
+SELECT quote(arrow_cast('it''s', 'Utf8View'));
+----
+'it\'s'
+
+query T
+SELECT quote(arrow_cast('', 'Utf8View'));
+----
+''
+
+# Multi-row batches, so the output builder is reused across rows
+
+query T
+SELECT quote(c) FROM VALUES ('a'), ('it''s'), (NULL), (''), ('x''''y') AS t(c);
+----
+'a'
+'it\'s'
+NULL
+''
+'x\'\'y'
+
+query T
+SELECT quote(arrow_cast(c, 'LargeUtf8')) FROM VALUES ('a'), ('it''s'), (NULL) AS t(c);
+----
+'a'
+'it\'s'
+NULL
+
+# Mixes values that are inlined in the views (<= 12 bytes) with values held in a
+# data buffer, since the two are measured differently for the capacity hint.
+query T
+SELECT quote(arrow_cast(c, 'Utf8View'))
+FROM VALUES ('a'), ('it''s'), (NULL), ('a value that is longer than twelve bytes and has an '' embedded quote') AS t(c);
+----
+'a'
+'it\'s'
+NULL
+'a value that is longer than twelve bytes and has an \' embedded quote'
diff --git a/datafusion/sqllogictest/test_files/spark/string/soundex.slt b/datafusion/sqllogictest/test_files/spark/string/soundex.slt
index ec85c4bd40..4a05d57471 100644
--- a/datafusion/sqllogictest/test_files/spark/string/soundex.slt
+++ b/datafusion/sqllogictest/test_files/spark/string/soundex.slt
@@ -199,3 +199,56 @@ query T
SELECT concat(soundex(' '), 'Spark')
----
Spark
+
+# LargeUtf8 and Utf8View inputs
+
+query T
+SELECT soundex(arrow_cast('Miller', 'LargeUtf8'));
+----
+M460
+
+query T
+SELECT soundex(arrow_cast('123', 'LargeUtf8'));
+----
+123
+
+query T
+SELECT soundex(arrow_cast('Miller', 'Utf8View'));
+----
+M460
+
+query T
+SELECT soundex(arrow_cast('123', 'Utf8View'));
+----
+123
+
+# Non-ASCII alphabetic first characters take the passthrough route
+
+query T
+SELECT soundex('Ñoño');
+----
+Ñoño
+
+query T
+SELECT soundex('Éclair');
+----
+Éclair
+
+# Multi-row batches, so the output builder is reused across rows
+
+query T
+SELECT soundex(c) FROM VALUES ('Miller'), (NULL), (''), ('Ñoño'), ('123'), ('Robert') AS t(c);
+----
+M460
+NULL
+(empty)
+Ñoño
+123
+R163
+
+query T
+SELECT soundex(arrow_cast(c, 'Utf8View')) FROM VALUES ('Miller'), (NULL), ('Ñoño') AS t(c);
+----
+M460
+NULL
+Ñoño
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org