diff --git a/Cargo.lock b/Cargo.lock index f3c07d27bc376..08ec2b54d579d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2617,6 +2617,7 @@ name = "datafusion-spark" version = "54.1.0" dependencies = [ "arrow", + "base64 0.23.0", "bigdecimal", "chrono", "crc32fast", diff --git a/datafusion/spark/Cargo.toml b/datafusion/spark/Cargo.toml index 93987b553f2f5..c3fd08f994a20 100644 --- a/datafusion/spark/Cargo.toml +++ b/datafusion/spark/Cargo.toml @@ -44,6 +44,7 @@ name = "datafusion_spark" [dependencies] arrow = { workspace = true } +base64 = "0.23" bigdecimal = { workspace = true } chrono = { workspace = true } crc32fast = "1.4" diff --git a/datafusion/spark/src/function/math/unhex.rs b/datafusion/spark/src/function/math/unhex.rs index 6739e6a15c582..941dc4330f015 100644 --- a/datafusion/spark/src/function/math/unhex.rs +++ b/datafusion/spark/src/function/math/unhex.rs @@ -147,7 +147,7 @@ where } /// Convert a single hex string to binary -fn unhex_scalar(s: &str) -> Option> { +pub(crate) fn unhex_scalar(s: &str) -> Option> { let mut buffer = Vec::with_capacity(s.len().div_ceil(2)); if unhex_common(s.as_bytes(), &mut buffer) { Some(buffer) diff --git a/datafusion/spark/src/function/string/mod.rs b/datafusion/spark/src/function/string/mod.rs index bc94c27732c91..0cadcdb91a0fb 100644 --- a/datafusion/spark/src/function/string/mod.rs +++ b/datafusion/spark/src/function/string/mod.rs @@ -32,6 +32,7 @@ pub mod quote; pub mod soundex; pub mod space; pub mod substring; +pub mod to_binary; use datafusion_expr::ScalarUDF; use datafusion_functions::make_udf_function; @@ -55,6 +56,8 @@ make_udf_function!(soundex::SparkSoundex, soundex); make_udf_function!(make_valid_utf8::SparkMakeValidUtf8, make_valid_utf8); make_udf_function!(is_valid_utf8::SparkIsValidUtf8, is_valid_utf8); make_udf_function!(quote::SparkQuote, quote); +make_udf_function!(to_binary::SparkToBinary, to_binary); +make_udf_function!(to_binary::SparkTryToBinary, try_to_binary); pub mod expr_fn { use datafusion_functions::export_functions; @@ -141,6 +144,16 @@ pub mod expr_fn { "Returns str enclosed by single quotes and each instance of single quote in it is preceded by a backslash", str )); + export_functions!(( + to_binary, + "Converts the input str to a binary value based on the supplied fmt, which must be a case-insensitive literal of 'hex', 'utf-8', 'utf8' or 'base64'. Defaults to 'hex'.", + str fmt + )); + export_functions!(( + try_to_binary, + "Like to_binary, but returns NULL instead of raising an error when the conversion cannot be performed.", + str fmt + )); } pub fn functions() -> Vec> { @@ -163,5 +176,7 @@ pub fn functions() -> Vec> { make_valid_utf8(), is_valid_utf8(), quote(), + to_binary(), + try_to_binary(), ] } diff --git a/datafusion/spark/src/function/string/to_binary.rs b/datafusion/spark/src/function/string/to_binary.rs new file mode 100644 index 0000000000000..db3b1c0bda345 --- /dev/null +++ b/datafusion/spark/src/function/string/to_binary.rs @@ -0,0 +1,494 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::Arc; + +use arrow::array::{Array, ArrayRef, AsArray, BinaryBuilder}; +use arrow::datatypes::DataType; +use base64::Engine as _; +use base64::engine::DecodePaddingMode; +use base64::engine::general_purpose::{GeneralPurpose, GeneralPurposeConfig}; +use datafusion_common::{Result, ScalarValue, exec_err, plan_err}; +use datafusion_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, + Volatility, +}; + +use crate::function::math::unhex::unhex_scalar; + +/// The binary formats accepted by `to_binary` / `try_to_binary`. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum BinaryFormat { + Hex, + Utf8, + Base64, +} + +impl BinaryFormat { + /// Spark matches the format case-insensitively against a fixed set. + fn parse(fmt: &str) -> Option { + match fmt.to_lowercase().as_str() { + "hex" => Some(Self::Hex), + "utf-8" | "utf8" => Some(Self::Utf8), + "base64" => Some(Self::Base64), + _ => None, + } + } +} + +/// Spark-compatible `to_binary` expression +/// +/// +/// `to_binary(str[, fmt])` converts `str` to binary using `fmt`, which must be a +/// case-insensitive literal of `hex`, `utf-8`, `utf8` or `base64`. `fmt` defaults +/// to `hex`. The result is NULL if any input is NULL. +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct SparkToBinary { + signature: Signature, +} + +impl Default for SparkToBinary { + fn default() -> Self { + Self::new() + } +} + +impl SparkToBinary { + pub fn new() -> Self { + Self { + signature: to_binary_signature(), + } + } +} + +impl ScalarUDFImpl for SparkToBinary { + fn name(&self) -> &str { + "to_binary" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Binary) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + to_binary_inner(self.name(), &args.args, true) + } +} + +/// Spark-compatible `try_to_binary` expression +/// +/// +/// Identical to [`SparkToBinary`] except that a value which cannot be converted, +/// and a format which is not recognised, both yield NULL instead of an error. +/// Spark expresses this as `TryEval(ToBinary(expr, fmt, nullOnInvalidFormat = true))`. +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct SparkTryToBinary { + signature: Signature, +} + +impl Default for SparkTryToBinary { + fn default() -> Self { + Self::new() + } +} + +impl SparkTryToBinary { + pub fn new() -> Self { + Self { + signature: to_binary_signature(), + } + } +} + +impl ScalarUDFImpl for SparkTryToBinary { + fn name(&self) -> &str { + "try_to_binary" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Binary) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + to_binary_inner(self.name(), &args.args, false) + } +} + +fn to_binary_signature() -> Signature { + let mut variants = Vec::with_capacity(12); + for str_type in [DataType::Utf8, DataType::Utf8View, DataType::LargeUtf8] { + variants.push(TypeSignature::Exact(vec![str_type.clone()])); + for fmt_type in [DataType::Utf8, DataType::Utf8View, DataType::LargeUtf8] { + variants.push(TypeSignature::Exact(vec![str_type.clone(), fmt_type])); + } + } + Signature::one_of(variants, Volatility::Immutable) +} + +/// Reads the format argument. Spark requires it to be foldable, so only a scalar +/// is accepted here; a NULL format makes the whole expression NULL. +fn read_format( + name: &str, + fmt: Option<&ColumnarValue>, + fail_on_error: bool, +) -> Result> { + let Some(fmt) = fmt else { + return Ok(Some(BinaryFormat::Hex)); // the default when `fmt` is omitted + }; + let ColumnarValue::Scalar(scalar) = fmt else { + return plan_err!("{name}: the fmt argument must be a constant, not a column"); + }; + let fmt = match scalar { + ScalarValue::Utf8(v) | ScalarValue::Utf8View(v) | ScalarValue::LargeUtf8(v) => v, + other => { + return plan_err!("{name}: the fmt argument must be a string, got {other:?}"); + } + }; + match fmt { + // A NULL format yields NULL rather than an error, in both variants. + None => Ok(None), + Some(fmt) => match BinaryFormat::parse(fmt) { + Some(parsed) => Ok(Some(parsed)), + // try_to_binary returns NULL for an unrecognised format + // (`nullOnInvalidFormat`); to_binary rejects it. + None if !fail_on_error => Ok(None), + None => plan_err!( + "{name}: invalid fmt '{fmt}', expected one of \ + 'hex', 'utf-8', 'utf8' or 'base64'" + ), + }, + } +} + +fn to_binary_inner( + name: &str, + args: &[ColumnarValue], + fail_on_error: bool, +) -> Result { + let (value, fmt) = match args { + [value] => (value, None), + [value, fmt] => (value, Some(fmt)), + _ => { + return exec_err!("{name} expects 1 or 2 arguments, got {}", args.len()); + } + }; + + let Some(format) = read_format(name, fmt, fail_on_error)? else { + // NULL or (for try_to_binary) unrecognised format: the whole result is NULL. + return Ok(ColumnarValue::Scalar(ScalarValue::Binary(None))); + }; + + match value { + ColumnarValue::Array(array) => { + let converted = convert_array(name, array, format, fail_on_error)?; + Ok(ColumnarValue::Array(converted)) + } + ColumnarValue::Scalar(scalar) => { + let value = match scalar { + ScalarValue::Utf8(v) + | ScalarValue::Utf8View(v) + | ScalarValue::LargeUtf8(v) => v.as_deref(), + other => { + return exec_err!( + "{name}: expected a string argument, got {other:?}" + ); + } + }; + let converted = match value { + None => None, + Some(v) => convert_one(name, v, format, fail_on_error)?, + }; + Ok(ColumnarValue::Scalar(ScalarValue::Binary(converted))) + } + } +} + +fn convert_array( + name: &str, + array: &ArrayRef, + format: BinaryFormat, + fail_on_error: bool, +) -> Result { + let values: Vec> = match array.data_type() { + DataType::Utf8 => array.as_string::().iter().collect(), + DataType::LargeUtf8 => array.as_string::().iter().collect(), + DataType::Utf8View => array.as_string_view().iter().collect(), + other => { + return exec_err!("{name}: expected a string argument, got {other}"); + } + }; + + let mut builder = BinaryBuilder::with_capacity(values.len(), array.len()); + for value in values { + match value { + None => builder.append_null(), + Some(v) => match convert_one(name, v, format, fail_on_error)? { + Some(bytes) => builder.append_value(&bytes), + None => builder.append_null(), + }, + } + } + Ok(Arc::new(builder.finish())) +} + +/// Converts one value. `Ok(None)` means "NULL", which only happens when +/// `fail_on_error` is false; otherwise an invalid value is an error. +fn convert_one( + name: &str, + value: &str, + format: BinaryFormat, + fail_on_error: bool, +) -> Result>> { + let converted = match format { + BinaryFormat::Utf8 => Some(value.as_bytes().to_vec()), + BinaryFormat::Hex => unhex_scalar(value), + BinaryFormat::Base64 => decode_base64(value), + }; + match converted { + Some(bytes) => Ok(Some(bytes)), + None if fail_on_error => exec_err!( + "{name}: cannot convert '{value}' to binary using format '{}'", + match format { + BinaryFormat::Hex => "hex", + BinaryFormat::Utf8 => "utf-8", + BinaryFormat::Base64 => "base64", + } + ), + None => Ok(None), + } +} + +/// The base64 engine Spark's decoding amounts to: the standard alphabet, the +/// padding optional, and the unused trailing bits of a short final group +/// ignored, as Java's MIME decoder does. +const SPARK_BASE64_DECODE: GeneralPurpose = GeneralPurpose::new( + &base64::alphabet::STANDARD, + GeneralPurposeConfig::new() + .with_decode_allow_trailing_bits(true) + .with_decode_padding_mode(DecodePaddingMode::Indifferent), +); + +fn is_base64_byte(byte: u8) -> bool { + byte.is_ascii_alphanumeric() || matches!(byte, b'+' | b'/' | b'=') +} + +/// Decodes a base64 string, returning `None` if it is not valid base64. +/// +/// Java's MIME decoder skips characters outside the alphabet, which the engine +/// has no setting for, so a failed decode is retried with those characters +/// removed. +fn decode_base64(value: &str) -> Option> { + let bytes = value.as_bytes(); + match SPARK_BASE64_DECODE.decode(bytes) { + Ok(decoded) => Some(decoded), + Err(_) if bytes.iter().any(|byte| !is_base64_byte(*byte)) => { + let filtered: Vec = bytes + .iter() + .copied() + .filter(|b| is_base64_byte(*b)) + .collect(); + SPARK_BASE64_DECODE.decode(filtered).ok() + } + Err(_) => None, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::StringArray; + use datafusion_common::internal_err; + + fn to_binary(value: &str, fmt: Option<&str>, fail: bool) -> Result>> { + let mut args = vec![ColumnarValue::Scalar(ScalarValue::Utf8(Some( + value.to_string(), + )))]; + if let Some(fmt) = fmt { + args.push(ColumnarValue::Scalar(ScalarValue::Utf8(Some( + fmt.to_string(), + )))); + } + let name = if fail { "to_binary" } else { "try_to_binary" }; + match to_binary_inner(name, &args, fail)? { + ColumnarValue::Scalar(ScalarValue::Binary(v)) => Ok(v), + other => internal_err!("unexpected result {other:?}"), + } + } + + #[test] + fn test_utf8_format() -> Result<()> { + assert_eq!( + to_binary("abc", Some("utf-8"), true)?, + Some(b"abc".to_vec()) + ); + assert_eq!(to_binary("abc", Some("utf8"), true)?, Some(b"abc".to_vec())); + assert_eq!( + to_binary("abc", Some("UTF-8"), true)?, + Some(b"abc".to_vec()) + ); + // multi-byte input round-trips as its UTF-8 bytes + assert_eq!(to_binary("é", Some("utf-8"), true)?, Some(vec![0xc3, 0xa9])); + Ok(()) + } + + #[test] + fn test_hex_is_the_default_format() -> Result<()> { + assert_eq!( + to_binary("537061726B", None, true)?, + Some(b"Spark".to_vec()) + ); + assert_eq!( + to_binary("537061726B", Some("hex"), true)?, + Some(b"Spark".to_vec()) + ); + // an odd number of digits is left-padded with '0' + assert_eq!(to_binary("F", None, true)?, Some(vec![0x0f])); + Ok(()) + } + + #[test] + fn test_base64_format() -> Result<()> { + assert_eq!( + to_binary("U3Bhcms=", Some("base64"), true)?, + Some(b"Spark".to_vec()) + ); + assert_eq!( + to_binary("YWJj", Some("base64"), true)?, + Some(b"abc".to_vec()) + ); + // whitespace between symbols is ignored + assert_eq!( + to_binary("U3Bh\ncms=", Some("base64"), true)?, + Some(b"Spark".to_vec()) + ); + // an unpadded final group is accepted + assert_eq!( + to_binary("U3Bhcms", Some("base64"), true)?, + Some(b"Spark".to_vec()) + ); + Ok(()) + } + + #[test] + fn test_empty_input() -> Result<()> { + for fmt in ["hex", "utf-8", "base64"] { + assert_eq!(to_binary("", Some(fmt), true)?, Some(vec![])); + } + Ok(()) + } + + #[test] + fn test_invalid_value_errors_or_nulls() -> Result<()> { + // to_binary raises error, try_to_binary returns NULL + assert!(to_binary("zz", Some("hex"), true).is_err()); + assert_eq!(to_binary("zz", Some("hex"), false)?, None); + + assert!(to_binary("a!", Some("base64"), true).is_err()); + assert_eq!(to_binary("a!", Some("base64"), false)?, None); + Ok(()) + } + + #[test] + fn test_invalid_format() -> Result<()> { + // to_binary rejects an unrecognised format, try_to_binary returns NULL + assert!(to_binary("abc", Some("invalidFormat"), true).is_err()); + assert_eq!(to_binary("abc", Some("invalidFormat"), false)?, None); + Ok(()) + } + + #[test] + fn test_null_inputs() -> Result<()> { + // a NULL value yields NULL + let args = vec![ColumnarValue::Scalar(ScalarValue::Utf8(None))]; + assert!(matches!( + to_binary_inner("to_binary", &args, true)?, + ColumnarValue::Scalar(ScalarValue::Binary(None)) + )); + + // a NULL format yields NULL, and is not an error even for to_binary + let args = vec![ + ColumnarValue::Scalar(ScalarValue::Utf8(Some("abc".into()))), + ColumnarValue::Scalar(ScalarValue::Utf8(None)), + ]; + assert!(matches!( + to_binary_inner("to_binary", &args, true)?, + ColumnarValue::Scalar(ScalarValue::Binary(None)) + )); + Ok(()) + } + + #[test] + fn test_column_input() -> Result<()> { + let array: ArrayRef = + Arc::new(StringArray::from(vec![Some("abc"), None, Some("déf")])); + let args = vec![ + ColumnarValue::Array(array), + ColumnarValue::Scalar(ScalarValue::Utf8(Some("utf-8".into()))), + ]; + let ColumnarValue::Array(result) = to_binary_inner("to_binary", &args, true)? + else { + unreachable!() + }; + let result = result.as_binary::(); + assert_eq!(result.value(0), b"abc"); + assert!(result.is_null(1)); + assert_eq!(result.value(2), "déf".as_bytes()); + Ok(()) + } + + #[test] + fn test_column_with_invalid_value() -> Result<()> { + let array: ArrayRef = Arc::new(StringArray::from(vec![Some("4142"), Some("zz")])); + let args = vec![ + ColumnarValue::Array(Arc::clone(&array)), + ColumnarValue::Scalar(ScalarValue::Utf8(Some("hex".into()))), + ]; + // to_binary fails the whole batch + assert!(to_binary_inner("to_binary", &args, true).is_err()); + + // try_to_binary nulls only the offending row + let ColumnarValue::Array(result) = + to_binary_inner("try_to_binary", &args, false)? + else { + unreachable!() + }; + let result = result.as_binary::(); + assert_eq!(result.value(0), b"AB"); + assert!(result.is_null(1)); + Ok(()) + } + + #[test] + fn test_base64_validation() { + assert_eq!(decode_base64("YWJj"), Some(b"abc".to_vec())); + assert_eq!(decode_base64(""), Some(vec![])); + assert_eq!(decode_base64("a!"), None); // invalid character + assert_eq!(decode_base64("YQ==="), None); // too much padding + assert_eq!(decode_base64("YQ=x"), None); // data after padding + assert_eq!(decode_base64("YWJjY"), None); // dangling symbol + assert_eq!(decode_base64("YWJjYQ=="), Some(b"abca".to_vec())); + assert_eq!(decode_base64("YWJj=="), None); // final group already complete + } +} diff --git a/datafusion/sqllogictest/test_files/spark/string/to_binary.slt b/datafusion/sqllogictest/test_files/spark/string/to_binary.slt index d8efa323f2c52..5a0ad2c23e5b9 100644 --- a/datafusion/sqllogictest/test_files/spark/string/to_binary.slt +++ b/datafusion/sqllogictest/test_files/spark/string/to_binary.slt @@ -23,5 +23,57 @@ ## Original Query: SELECT to_binary('abc', 'utf-8'); ## PySpark 3.5.5 Result: {'to_binary(abc, utf-8)': bytearray(b'abc'), 'typeof(to_binary(abc, utf-8))': 'binary', 'typeof(abc)': 'string', 'typeof(utf-8)': 'string'} -#query -#SELECT to_binary('abc'::string, 'utf-8'::string); +query ? +SELECT to_binary('abc'::string, 'utf-8'::string); +---- +616263 + +# the format is matched case-insensitively, and utf8 is an accepted spelling +query ? +SELECT to_binary('abc'::string, 'UTF8'::string); +---- +616263 + +# hex is the default when fmt is omitted +query ? +SELECT to_binary('537061726b'::string); +---- +537061726b + +query ? +SELECT to_binary('537061726b'::string, 'hex'::string); +---- +537061726b + +# an odd number of hex digits is left-padded with a leading zero +query ? +SELECT to_binary('f'::string, 'hex'::string); +---- +0f + +query ? +SELECT to_binary('U3Bhcms='::string, 'base64'::string); +---- +537061726b + +# NULL input, and a NULL format, both give NULL +query ? +SELECT to_binary(NULL::string, 'hex'::string); +---- +NULL + +query ? +SELECT to_binary('abc'::string, NULL::string); +---- +NULL + +# to_binary raises on a value that cannot be converted +statement error +SELECT to_binary('zz'::string, 'hex'::string); + +statement error +SELECT to_binary('a!'::string, 'base64'::string); + +# ...and on a format it does not recognise +statement error +SELECT to_binary('abc'::string, 'invalidFormat'::string); diff --git a/datafusion/sqllogictest/test_files/spark/string/try_to_binary.slt b/datafusion/sqllogictest/test_files/spark/string/try_to_binary.slt index 211520be1e48b..74ff1041936d2 100644 --- a/datafusion/sqllogictest/test_files/spark/string/try_to_binary.slt +++ b/datafusion/sqllogictest/test_files/spark/string/try_to_binary.slt @@ -23,15 +23,43 @@ ## Original Query: SELECT try_to_binary('abc', 'utf-8'); ## PySpark 3.5.5 Result: {'try_to_binary(abc, utf-8)': bytearray(b'abc'), 'typeof(try_to_binary(abc, utf-8))': 'binary', 'typeof(abc)': 'string', 'typeof(utf-8)': 'string'} -#query -#SELECT try_to_binary('abc'::string, 'utf-8'::string); +query ? +SELECT try_to_binary('abc'::string, 'utf-8'::string); +---- +616263 ## Original Query: select try_to_binary('a!', 'base64'); ## PySpark 3.5.5 Result: {'try_to_binary(a!, base64)': None, 'typeof(try_to_binary(a!, base64))': 'binary', 'typeof(a!)': 'string', 'typeof(base64)': 'string'} -#query -#SELECT try_to_binary('a!'::string, 'base64'::string); +query ? +SELECT try_to_binary('a!'::string, 'base64'::string); +---- +NULL ## Original Query: select try_to_binary('abc', 'invalidFormat'); ## PySpark 3.5.5 Result: {'try_to_binary(abc, invalidFormat)': None, 'typeof(try_to_binary(abc, invalidFormat))': 'binary', 'typeof(abc)': 'string', 'typeof(invalidFormat)': 'string'} -#query -#SELECT try_to_binary('abc'::string, 'invalidFormat'::string); +query ? +SELECT try_to_binary('abc'::string, 'invalidFormat'::string); +---- +NULL + +# where the conversion succeeds, try_to_binary matches to_binary +query ? +SELECT try_to_binary('537061726b'::string); +---- +537061726b + +query ? +SELECT try_to_binary('U3Bhcms='::string, 'base64'::string); +---- +537061726b + +# an invalid value is NULL rather than an error +query ? +SELECT try_to_binary('zz'::string, 'hex'::string); +---- +NULL + +query ? +SELECT try_to_binary(NULL::string, 'hex'::string); +---- +NULL