diff --git a/datafusion/spark/src/function/string/quote.rs b/datafusion/spark/src/function/string/quote.rs index 55d827fb7a280..b4fd5e3827769 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 { 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)) + // 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::(str_array.iter(), 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) + // `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 = '\''; -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>>( + input: I, + data_capacity: usize, +) -> ArrayRef { + let len = input.size_hint().0; + let mut builder = + GenericStringBuilder::::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 GenericStringBuilder, 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> { + as_generic_string_array::(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::(&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 1fef0d5384821..8313a15b8108e 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 { } 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, O: OffsetSizeTrait, I: Iterator>>( + input: I, +) -> ArrayRef { + let len = input.size_hint().0; + // A soundex code is always exactly 4 ASCII characters. + let mut builder = GenericStringBuilder::::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 +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(builder: &mut GenericStringBuilder, 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> { + as_generic_string_array::(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::(&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 b5ef0f84e60d2..e63bf196d7a38 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 ec85c4bd40b24..4a05d57471d3a 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