From 5205d7045815d49d704310d6edac2ce528342e49 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 25 Jul 2026 08:07:34 -0600 Subject: [PATCH 1/2] perf: remove per-row String allocation from Spark url_encode url_encode built a String for every row via byte_serialize(..).collect::() and collected the results into a StringArray, so each row paid for its own allocation. The percent-encoded form is now assembled in a single scratch buffer reused across rows and appended to a pre-sized builder. The now-unused UrlEncode::encode helper, which returned a Result that was never an Err, is removed. -56% at 1024 rows and -48% at 8192 against the benchmark added in #23882. --- .../spark/src/function/url/url_encode.rs | 71 +++++++++++-------- 1 file changed, 41 insertions(+), 30 deletions(-) diff --git a/datafusion/spark/src/function/url/url_encode.rs b/datafusion/spark/src/function/url/url_encode.rs index 1ad2a111851ee..b54f66561f2df 100644 --- a/datafusion/spark/src/function/url/url_encode.rs +++ b/datafusion/spark/src/function/url/url_encode.rs @@ -17,7 +17,9 @@ use std::sync::Arc; -use arrow::array::{ArrayRef, LargeStringArray, StringArray, StringViewArray}; +use arrow::array::{ + Array, ArrayRef, LargeStringBuilder, StringBuilder, StringViewBuilder, +}; use arrow::datatypes::DataType; use datafusion_common::cast::{ as_large_string_array, as_string_array, as_string_view_array, @@ -46,20 +48,6 @@ impl UrlEncode { signature: Signature::string(1, Volatility::Immutable), } } - - /// Encode a string to application/x-www-form-urlencoded format. - /// - /// # Arguments - /// - /// * `value` - The string to encode - /// - /// # Returns - /// - /// * `Ok(String)` - The encoded string - /// - fn encode(value: &str) -> Result { - Ok(byte_serialize(value.as_bytes()).collect::()) - } } impl ScalarUDFImpl for UrlEncode { @@ -105,22 +93,45 @@ fn spark_url_encode(args: &[ArrayRef]) -> Result { return exec_err!("`url_encode` expects 1 argument"); } + // The percent-encoded form of each value is assembled in a single scratch buffer + // reused across rows, rather than allocating a `String` per row. + macro_rules! encode_all { + ($array:expr, $builder:expr) => {{ + let array = $array; + let mut builder = $builder; + let mut encoded = String::new(); + for value in array.iter() { + match value { + Some(value) => { + encoded.clear(); + encoded.extend(byte_serialize(value.as_bytes())); + builder.append_value(&encoded); + } + None => builder.append_null(), + } + } + Ok(Arc::new(builder.finish()) as ArrayRef) + }}; + } + match &args[0].data_type() { - DataType::Utf8 => as_string_array(&args[0])? - .iter() - .map(|x| x.map(UrlEncode::encode).transpose()) - .collect::>() - .map(|array| Arc::new(array) as ArrayRef), - DataType::LargeUtf8 => as_large_string_array(&args[0])? - .iter() - .map(|x| x.map(UrlEncode::encode).transpose()) - .collect::>() - .map(|array| Arc::new(array) as ArrayRef), - DataType::Utf8View => as_string_view_array(&args[0])? - .iter() - .map(|x| x.map(UrlEncode::encode).transpose()) - .collect::>() - .map(|array| Arc::new(array) as ArrayRef), + DataType::Utf8 => { + let array = as_string_array(&args[0])?; + let builder = + StringBuilder::with_capacity(array.len(), array.value_data().len()); + encode_all!(array, builder) + } + DataType::LargeUtf8 => { + let array = as_large_string_array(&args[0])?; + let builder = + LargeStringBuilder::with_capacity(array.len(), array.value_data().len()); + encode_all!(array, builder) + } + DataType::Utf8View => { + let array = as_string_view_array(&args[0])?; + let builder = StringViewBuilder::with_capacity(array.len()); + encode_all!(array, builder) + } other => exec_err!("`url_encode`: Expr must be STRING, got {other:?}"), } } From 3b3a7487711689e877c20696150508daca0f8e5b Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 25 Jul 2026 09:08:43 -0600 Subject: [PATCH 2/2] perf: avoid per-row String in url_decode, and build parse_url's null key array directly url_decode ended with .map(|parsed| parsed.into_owned()). Both replace_plus and decode_utf8 borrow their input when there is nothing to rewrite, so that allocated a String for every row even when the value needed no decoding at all. decode now returns a Cow and the array paths append it to a pre-sized builder. To let the Cow survive, spark_handled_url_decode's error handler changes from a closure over Result> to an OnDecodeError enum. Its only caller, try_url_decode, passed a closure that mapped Err to Ok(None), which is exactly OnDecodeError::Null. parse_url built its all-null key array for the 2-argument form by calling append_null() once per row on an uncapacitated builder; new_null_array does it in one allocation. url_decode -29% where nothing needs unescaping, -4% where it does, against the benchmark added in #23882. --- .../spark/src/function/url/parse_url.rs | 22 ++-- .../spark/src/function/url/try_url_decode.rs | 9 +- .../spark/src/function/url/url_decode.rs | 100 ++++++++++++------ 3 files changed, 84 insertions(+), 47 deletions(-) diff --git a/datafusion/spark/src/function/url/parse_url.rs b/datafusion/spark/src/function/url/parse_url.rs index 18f0bb1e0d78b..9ceed8b155bbd 100644 --- a/datafusion/spark/src/function/url/parse_url.rs +++ b/datafusion/spark/src/function/url/parse_url.rs @@ -18,8 +18,8 @@ use std::sync::Arc; use arrow::array::{ - Array, ArrayRef, GenericStringBuilder, LargeStringArray, StringArray, - StringArrayType, StringViewArray, + Array, ArrayRef, AsArray, LargeStringArray, StringArray, StringArrayType, + StringViewArray, new_null_array, }; use arrow::datatypes::DataType; use datafusion_common::cast::{ @@ -272,20 +272,18 @@ pub fn spark_handled_parse_url( ), } } else { - // The 'key' argument is omitted, assume all values are null - // Create 'null' string array for 'key' argument - let mut builder: GenericStringBuilder = GenericStringBuilder::new(); - for _ in 0..args[0].len() { - builder.append_null(); - } - let key = builder.finish(); + // The 'key' argument is omitted, assume all values are null. + // `new_null_array` allocates the null array outright, rather than + // appending one null per row through a builder. + let key_array = new_null_array(&DataType::Utf8, args[0].len()); + let key = key_array.as_string::(); match (url.data_type(), part.data_type()) { (DataType::Utf8, DataType::Utf8) => { process_parse_url::<_, _, _, StringArray>( as_string_array(url)?, as_string_array(part)?, - &key, + key, handler_err, false, ) @@ -294,7 +292,7 @@ pub fn spark_handled_parse_url( process_parse_url::<_, _, _, StringViewArray>( as_string_view_array(url)?, as_string_view_array(part)?, - &key, + key, handler_err, false, ) @@ -303,7 +301,7 @@ pub fn spark_handled_parse_url( process_parse_url::<_, _, _, LargeStringArray>( as_large_string_array(url)?, as_large_string_array(part)?, - &key, + key, handler_err, false, ) diff --git a/datafusion/spark/src/function/url/try_url_decode.rs b/datafusion/spark/src/function/url/try_url_decode.rs index 78968288fc2f5..1acbd5e13b988 100644 --- a/datafusion/spark/src/function/url/try_url_decode.rs +++ b/datafusion/spark/src/function/url/try_url_decode.rs @@ -24,7 +24,9 @@ use datafusion_expr::{ }; use datafusion_functions::utils::make_scalar_function; -use crate::function::url::url_decode::{UrlDecode, spark_handled_url_decode}; +use crate::function::url::url_decode::{ + OnDecodeError, UrlDecode, spark_handled_url_decode, +}; #[derive(Debug, PartialEq, Eq, Hash)] pub struct TryUrlDecode { @@ -67,10 +69,7 @@ impl ScalarUDFImpl for TryUrlDecode { } fn spark_try_url_decode(args: &[ArrayRef]) -> Result { - spark_handled_url_decode(args, |x| match x { - Err(_) => Ok(None), - result => result, - }) + spark_handled_url_decode(args, OnDecodeError::Null) } #[cfg(test)] diff --git a/datafusion/spark/src/function/url/url_decode.rs b/datafusion/spark/src/function/url/url_decode.rs index 0966cc380e497..879c83049d0cf 100644 --- a/datafusion/spark/src/function/url/url_decode.rs +++ b/datafusion/spark/src/function/url/url_decode.rs @@ -18,7 +18,9 @@ use std::borrow::Cow; use std::sync::Arc; -use arrow::array::{ArrayRef, LargeStringArray, StringArray, StringViewArray}; +use arrow::array::{ + Array, ArrayRef, LargeStringBuilder, StringBuilder, StringViewBuilder, +}; use arrow::datatypes::DataType; use datafusion_common::cast::{ as_large_string_array, as_string_array, as_string_view_array, @@ -61,18 +63,25 @@ impl UrlDecode { /// /// # Returns /// - /// * `Ok(String)` - The decoded string + /// * `Ok(Cow)` - The decoded string, borrowed from `value` when there + /// was nothing to rewrite and owned otherwise /// * `Err(DataFusionError)` - If the input is malformed or contains invalid UTF-8 - /// - fn decode(value: &str) -> Result { + fn decode(value: &str) -> Result> { // Check if the string has valid percent encoding Self::validate_percent_encoding(value)?; - let replaced = Self::replace_plus(value.as_bytes()); - percent_decode(&replaced) - .decode_utf8() - .map_err(|e| exec_datafusion_err!("Invalid UTF-8 sequence: {e}")) - .map(|parsed| parsed.into_owned()) + match Self::replace_plus(value.as_bytes()) { + // No '+' was rewritten, so the decode can borrow from `value` itself. + Cow::Borrowed(bytes) => percent_decode(bytes) + .decode_utf8() + .map_err(|e| exec_datafusion_err!("Invalid UTF-8 sequence: {e}")), + // Rewriting '+' already allocated, so owning the decoded form here + // costs nothing beyond what has been spent. + Cow::Owned(bytes) => percent_decode(&bytes) + .decode_utf8() + .map(|decoded| Cow::Owned(decoded.into_owned())) + .map_err(|e| exec_datafusion_err!("Invalid UTF-8 sequence: {e}")), + } } /// Replace b'+' with b' ' @@ -155,6 +164,15 @@ impl ScalarUDFImpl for UrlDecode { } } +/// How [`spark_handled_url_decode`] reacts to a malformed input value. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum OnDecodeError { + /// Propagate the error, as `url_decode` does. + Fail, + /// Return NULL for that row, as `try_url_decode` does. + Null, +} + /// Core implementation of URL decoding function. /// /// # Arguments @@ -165,38 +183,59 @@ impl ScalarUDFImpl for UrlDecode { /// /// * `Ok(ArrayRef)` - A new array of the same type containing decoded strings /// * `Err(DataFusionError)` - If validation fails or invalid arguments are provided -/// fn spark_url_decode(args: &[ArrayRef]) -> Result { - spark_handled_url_decode(args, |x| x) + spark_handled_url_decode(args, OnDecodeError::Fail) } pub fn spark_handled_url_decode( args: &[ArrayRef], - err_handle_fn: impl Fn(Result>) -> Result>, + on_error: OnDecodeError, ) -> Result { if args.len() != 1 { return exec_err!("`url_decode` expects 1 argument"); } + // Decoded values go straight into the builder, so a row that needs no + // unescaping is copied once rather than materialised as its own `String`. + macro_rules! decode_all { + ($array:expr, $builder:expr) => {{ + let array = $array; + let mut builder = $builder; + for value in array.iter() { + let Some(value) = value else { + builder.append_null(); + continue; + }; + match UrlDecode::decode(value) { + Ok(decoded) => builder.append_value(&decoded), + Err(e) => match on_error { + OnDecodeError::Fail => return Err(e), + OnDecodeError::Null => builder.append_null(), + }, + } + } + Ok(Arc::new(builder.finish()) as ArrayRef) + }}; + } + match &args[0].data_type() { - DataType::Utf8 => as_string_array(&args[0])? - .iter() - .map(|x| x.map(UrlDecode::decode).transpose()) - .map(&err_handle_fn) - .collect::>() - .map(|array| Arc::new(array) as ArrayRef), - DataType::LargeUtf8 => as_large_string_array(&args[0])? - .iter() - .map(|x| x.map(UrlDecode::decode).transpose()) - .map(&err_handle_fn) - .collect::>() - .map(|array| Arc::new(array) as ArrayRef), - DataType::Utf8View => as_string_view_array(&args[0])? - .iter() - .map(|x| x.map(UrlDecode::decode).transpose()) - .map(&err_handle_fn) - .collect::>() - .map(|array| Arc::new(array) as ArrayRef), + DataType::Utf8 => { + let array = as_string_array(&args[0])?; + let builder = + StringBuilder::with_capacity(array.len(), array.value_data().len()); + decode_all!(array, builder) + } + DataType::LargeUtf8 => { + let array = as_large_string_array(&args[0])?; + let builder = + LargeStringBuilder::with_capacity(array.len(), array.value_data().len()); + decode_all!(array, builder) + } + DataType::Utf8View => { + let array = as_string_view_array(&args[0])?; + let builder = StringViewBuilder::with_capacity(array.len()); + decode_all!(array, builder) + } other => exec_err!("`url_decode`: Expr must be STRING, got {other:?}"), } } @@ -205,6 +244,7 @@ pub fn spark_handled_url_decode( mod tests { use super::*; + use arrow::array::StringArray; #[test] fn test_decode() -> Result<()> {