diff --git a/datafusion/common/src/utils/hex.rs b/datafusion/common/src/utils/hex.rs index 872d54f40c6f7..3b08697469077 100644 --- a/datafusion/common/src/utils/hex.rs +++ b/datafusion/common/src/utils/hex.rs @@ -161,6 +161,9 @@ 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. +#[deprecated(note = "use encode_bytes_into or encode_bytes_to_slice instead")] +/// /// # Example /// /// ``` @@ -237,6 +240,7 @@ fn write_digits(v: u64, case: HexCase, buf: &mut [u8; 16]) -> usize { } #[cfg(test)] +#[expect(deprecated)] mod tests { use super::*; diff --git a/datafusion/functions/src/crypto/md5.rs b/datafusion/functions/src/crypto/md5.rs index b1206d2e423cc..c9aba15e6ef22 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 { 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 850e312abdb40..1f809b4e5da40 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_into, encode_bytes_to_slice}, take_function_args, }, }; @@ -373,7 +373,11 @@ 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 => { + let mut out = Vec::with_capacity(value.len() * 2); + encode_bytes_into(value, HexCase::Lower, &mut out); + unsafe { String::from_utf8_unchecked(out) } + } } } diff --git a/datafusion/spark/src/function/hash/sha1.rs b/datafusion/spark/src/function/hash/sha1.rs index 05a224f33f25a..64c52a7108283 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>) -> ArrayRef { diff --git a/datafusion/spark/src/function/hash/sha2.rs b/datafusion/spark/src/function/hash/sha2.rs index 541df2957669e..53c1ff9465781 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>, { - 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::(); - 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) } diff --git a/datafusion/spark/src/function/string/quote.rs b/datafusion/spark/src/function/string/quote.rs index 39ad8bf841764..55d827fb7a280 100644 --- a/datafusion/spark/src/function/string/quote.rs +++ b/datafusion/spark/src/function/string/quote.rs @@ -17,12 +17,14 @@ use arrow::array::{ArrayRef, OffsetSizeTrait, StringArray}; use arrow::datatypes::DataType; -use datafusion::logical_expr::{Coercion, ColumnarValue, Signature, TypeSignatureClass}; use datafusion_common::cast::{as_generic_string_array, as_string_view_array}; use datafusion_common::types::{NativeType, logical_string}; use datafusion_common::utils::take_function_args; use datafusion_common::{Result, exec_err}; -use datafusion_expr::{ScalarFunctionArgs, ScalarUDFImpl, Volatility}; +use datafusion_expr::{ + Coercion, ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, + TypeSignatureClass, Volatility, +}; use datafusion_functions::utils::make_scalar_function; use std::sync::Arc;