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-24057-f3ef45a53c9b59fe3e4802e08d9f878dbead9655 in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit a10d19374ffcb0f24253f819a9e254336561e0bb Author: RIchard Baah <[email protected]> AuthorDate: Sat Aug 15 11:06:08 2026 +0000 Perf: remove extra allocations for hex operations (#24057) ## Which issue does this PR close? <!-- We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes #123` indicates that this PR will close issue #123. --> - Closes #23810. ## Rationale for this change The digest functions hex-encode their output by allocating one String per row and then copying each of those into the output array (#23810) <!-- Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed. Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes. --> ## What changes are included in this PR? - This PR reduces per-row allocation and string-conversion overhead in the Spark crypto/hash paths by switching to the buffer-based hex encoder APIs. - Updates call sites to re-use vector buffers. <!-- There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR. --> ## Are these changes tested? existing test cover behavior. <!-- We typically require tests for all PRs in order to: 1. Prevent the code from being accidentally broken by subsequent changes 2. Serve as another way to document the expected behavior of the code If tests are not included in your PR, please explain why (for example, are they covered by existing tests)? --> ## Are there any user-facing changes? yes, `encode_bytes()` is being deprecated. This was apart of the public API <!-- If there are user-facing changes then we may require documentation to be updated before approving the PR. --> <!-- If there are any breaking changes to public APIs, please add the `api change` label. --> --------- Co-authored-by: rich-T-kid <[email protected]> --- datafusion/common/src/utils/hex.rs | 2 + datafusion/functions/src/crypto/md5.rs | 43 +++++++--- datafusion/functions/src/encoding/inner.rs | 4 +- datafusion/spark/src/function/hash/sha1.rs | 7 +- datafusion/spark/src/function/hash/sha2.rs | 121 +++++++++++++++++++++-------- 5 files changed, 131 insertions(+), 46 deletions(-) diff --git a/datafusion/common/src/utils/hex.rs b/datafusion/common/src/utils/hex.rs index 872d54f40c..672cd8463f 100644 --- a/datafusion/common/src/utils/hex.rs +++ b/datafusion/common/src/utils/hex.rs @@ -161,6 +161,8 @@ pub fn encode_bytes_to_slice(bytes: &[u8], case: HexCase, out: &mut [u8]) -> Res /// Returns the hex encoding of `bytes` as an owned `String`. /// +/// Prefer [`encode_bytes_into`] when you already have a reusable output buffer. +/// /// # Example /// /// ``` diff --git a/datafusion/functions/src/crypto/md5.rs b/datafusion/functions/src/crypto/md5.rs index b1206d2e42..c9aba15e6e 100644 --- a/datafusion/functions/src/crypto/md5.rs +++ b/datafusion/functions/src/crypto/md5.rs @@ -15,13 +15,16 @@ // specific language governing permissions and limitations // under the License. -use arrow::{array::StringViewArray, datatypes::DataType}; +use arrow::{ + array::{Array, BinaryViewBuilder}, + datatypes::DataType, +}; use datafusion_common::{ Result, ScalarValue, cast::as_binary_array, internal_err, types::{logical_binary, logical_string}, - utils::hex::{HexCase, encode_bytes}, + utils::hex::{HexCase, encode_bytes_into}, utils::take_function_args, }; use datafusion_expr::{ @@ -107,15 +110,35 @@ fn md5(args: &[ColumnarValue]) -> Result<ColumnarValue> { Ok(match value { ColumnarValue::Array(array) => { let binary_array = as_binary_array(&array)?; - let string_array: StringViewArray = binary_array - .iter() - .map(|opt| opt.map(|b| encode_bytes(b, HexCase::Lower))) - .collect(); - ColumnarValue::Array(Arc::new(string_array)) + let mut byte_builder = BinaryViewBuilder::with_capacity(binary_array.len()); + let mut hex_bytes = Vec::with_capacity(32); + + for i in 0..binary_array.len() { + if binary_array.is_null(i) { + byte_builder.append_null(); + continue; + } + + hex_bytes.clear(); + let digest = binary_array.value(i); + encode_bytes_into(digest, HexCase::Lower, &mut hex_bytes); + byte_builder.append_value(&hex_bytes); + } + + let str_array = unsafe { + // Safe: `encode_bytes_into` only writes ASCII hex digits, so the bytes are valid UTF-8. + byte_builder.finish().to_string_view_unchecked() + }; + ColumnarValue::Array(Arc::new(str_array)) + } + ColumnarValue::Scalar(ScalarValue::Binary(opt)) => { + ColumnarValue::Scalar(ScalarValue::Utf8View(opt.map(|b| { + let mut hex_bytes = Vec::with_capacity(b.len() * 2); + encode_bytes_into(&b, HexCase::Lower, &mut hex_bytes); + // Safe: `encode_bytes_into` only writes ASCII hex digits, so the bytes are valid UTF-8. + unsafe { String::from_utf8_unchecked(hex_bytes) } + }))) } - ColumnarValue::Scalar(ScalarValue::Binary(opt)) => ColumnarValue::Scalar( - ScalarValue::Utf8View(opt.map(|b| encode_bytes(&b, HexCase::Lower))), - ), _ => return internal_err!("Impossibly got invalid results from digest"), }) } diff --git a/datafusion/functions/src/encoding/inner.rs b/datafusion/functions/src/encoding/inner.rs index 850e312abd..8b57033fa0 100644 --- a/datafusion/functions/src/encoding/inner.rs +++ b/datafusion/functions/src/encoding/inner.rs @@ -34,7 +34,7 @@ use datafusion_common::{ not_impl_err, plan_err, types::{NativeType, logical_string}, utils::{ - hex::{HexCase, encode_bytes as encode_hex, encode_bytes_to_slice}, + hex::{HexCase, encode_bytes, encode_bytes_to_slice}, take_function_args, }, }; @@ -373,7 +373,7 @@ impl Encoding { match self { Self::Base64 => BASE64_ENGINE.encode(value), Self::Base64Padded => BASE64_ENGINE_PADDED.encode(value), - Self::Hex => encode_hex(value, HexCase::Lower), + Self::Hex => encode_bytes(value, HexCase::Lower), } } diff --git a/datafusion/spark/src/function/hash/sha1.rs b/datafusion/spark/src/function/hash/sha1.rs index 05a224f33f..64c52a7108 100644 --- a/datafusion/spark/src/function/hash/sha1.rs +++ b/datafusion/spark/src/function/hash/sha1.rs @@ -24,7 +24,7 @@ use datafusion_common::cast::{ as_large_binary_array, }; use datafusion_common::types::{NativeType, logical_string}; -use datafusion_common::utils::hex::{HexCase, encode_bytes}; +use datafusion_common::utils::hex::{HexCase, encode_bytes_into}; use datafusion_common::utils::take_function_args; use datafusion_common::{Result, internal_err}; use datafusion_expr::{ @@ -92,7 +92,10 @@ impl ScalarUDFImpl for SparkSha1 { #[inline] fn spark_sha1_digest(value: &[u8]) -> String { - encode_bytes(&Sha1::digest(value), HexCase::Lower) + let mut out = Vec::with_capacity(40); + // Safe: `encode_bytes_into` only writes ASCII hex digits, which are valid UTF-8. + encode_bytes_into(&Sha1::digest(value), HexCase::Lower, &mut out); + unsafe { String::from_utf8_unchecked(out) } } fn spark_sha1_impl<'a>(input: impl Iterator<Item = Option<&'a [u8]>>) -> ArrayRef { diff --git a/datafusion/spark/src/function/hash/sha2.rs b/datafusion/spark/src/function/hash/sha2.rs index 541df29576..53c1ff9465 100644 --- a/datafusion/spark/src/function/hash/sha2.rs +++ b/datafusion/spark/src/function/hash/sha2.rs @@ -15,12 +15,14 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::{ArrayRef, AsArray, BinaryArrayType, Int32Array, StringArray}; +use arrow::array::{ + ArrayRef, AsArray, BinaryArrayType, BinaryBuilder, Int32Array, StringArray, +}; use arrow::datatypes::{DataType, Int32Type}; use datafusion_common::types::{ NativeType, logical_binary, logical_int32, logical_string, }; -use datafusion_common::utils::hex::{HexCase, encode_bytes}; +use datafusion_common::utils::hex::{HexCase, encode_bytes_into}; use datafusion_common::utils::take_function_args; use datafusion_common::{Result, ScalarValue, internal_err}; use datafusion_expr::{ @@ -113,22 +115,58 @@ impl ScalarUDFImpl for SparkSha2 { 224 => { let mut digest = sha2::Sha224::default(); digest.update(bytes); - Some(encode_bytes(&digest.finalize(), HexCase::Lower)) + let mut hex_bytes = Vec::with_capacity(56); + encode_bytes_into( + &digest.finalize(), + HexCase::Lower, + &mut hex_bytes, + ); + Some( + String::from_utf8(hex_bytes) + .expect("ASCII hex is valid UTF-8"), + ) } 0 | 256 => { let mut digest = sha2::Sha256::default(); digest.update(bytes); - Some(encode_bytes(&digest.finalize(), HexCase::Lower)) + let mut hex_bytes = Vec::with_capacity(64); + encode_bytes_into( + &digest.finalize(), + HexCase::Lower, + &mut hex_bytes, + ); + Some( + String::from_utf8(hex_bytes) + .expect("ASCII hex is valid UTF-8"), + ) } 384 => { let mut digest = sha2::Sha384::default(); digest.update(bytes); - Some(encode_bytes(&digest.finalize(), HexCase::Lower)) + let mut hex_bytes = Vec::with_capacity(96); + encode_bytes_into( + &digest.finalize(), + HexCase::Lower, + &mut hex_bytes, + ); + Some( + String::from_utf8(hex_bytes) + .expect("ASCII hex is valid UTF-8"), + ) } 512 => { let mut digest = sha2::Sha512::default(); digest.update(bytes); - Some(encode_bytes(&digest.finalize(), HexCase::Lower)) + let mut hex_bytes = Vec::with_capacity(128); + encode_bytes_into( + &digest.finalize(), + HexCase::Lower, + &mut hex_bytes, + ); + Some( + String::from_utf8(hex_bytes) + .expect("ASCII hex is valid UTF-8"), + ) } _ => None, }; @@ -216,33 +254,52 @@ where BinaryArrType: BinaryArrayType<'a>, I: Iterator<Item = Option<i32>>, { - let array = values + let mut byte_builder = BinaryBuilder::with_capacity(values.len(), values.len() * 2); + let mut hex_bytes = Vec::with_capacity(128); + + values .iter() .zip(bit_lengths) - .map(|(value, bit_length)| match (value, bit_length) { - (Some(value), Some(224)) => { - let mut digest = sha2::Sha224::default(); - digest.update(value); - Some(encode_bytes(&digest.finalize(), HexCase::Lower)) - } - (Some(value), Some(0 | 256)) => { - let mut digest = sha2::Sha256::default(); - digest.update(value); - Some(encode_bytes(&digest.finalize(), HexCase::Lower)) - } - (Some(value), Some(384)) => { - let mut digest = sha2::Sha384::default(); - digest.update(value); - Some(encode_bytes(&digest.finalize(), HexCase::Lower)) - } - (Some(value), Some(512)) => { - let mut digest = sha2::Sha512::default(); - digest.update(value); - Some(encode_bytes(&digest.finalize(), HexCase::Lower)) + .for_each(|(value, bit_length)| { + match (value, bit_length) { + (Some(value), Some(224)) => { + let mut digest = sha2::Sha224::default(); + digest.update(value); + hex_bytes.clear(); + encode_bytes_into(&digest.finalize(), HexCase::Lower, &mut hex_bytes); + byte_builder.append_value(&hex_bytes); + } + (Some(value), Some(0 | 256)) => { + let mut digest = sha2::Sha256::default(); + digest.update(value); + hex_bytes.clear(); + encode_bytes_into(&digest.finalize(), HexCase::Lower, &mut hex_bytes); + byte_builder.append_value(&hex_bytes); + } + (Some(value), Some(384)) => { + let mut digest = sha2::Sha384::default(); + digest.update(value); + hex_bytes.clear(); + encode_bytes_into(&digest.finalize(), HexCase::Lower, &mut hex_bytes); + byte_builder.append_value(&hex_bytes); + } + (Some(value), Some(512)) => { + let mut digest = sha2::Sha512::default(); + digest.update(value); + hex_bytes.clear(); + encode_bytes_into(&digest.finalize(), HexCase::Lower, &mut hex_bytes); + byte_builder.append_value(&hex_bytes); + } + // Unknown bit-lengths go to null, same as in Spark + _ => byte_builder.append_null(), } - // Unknown bit-lengths go to null, same as in Spark - _ => None, - }) - .collect::<StringArray>(); - Arc::new(array) + }); + + let str_array = unsafe { + let binary_array = byte_builder.finish(); + let (offsets, values, nulls) = binary_array.into_parts(); + // Safe: `encode_bytes_into` only writes ASCII hex digits, so the bytes are valid UTF-8. + StringArray::new_unchecked(offsets, values, nulls) + }; + Arc::new(str_array) } --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
