diff --git a/Cargo.lock b/Cargo.lock index 9deb9f4f8e089..0ab26a09c4f90 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2620,6 +2620,7 @@ name = "datafusion-spark" version = "54.1.0" dependencies = [ "arrow", + "base64 0.23.1", "bigdecimal", "chrono", "crc32fast", diff --git a/datafusion/spark/Cargo.toml b/datafusion/spark/Cargo.toml index 01fd9ec9b3f5b..f40708863d108 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..03ff8e5e58e09 --- /dev/null +++ b/datafusion/spark/src/function/string/to_binary.rs @@ -0,0 +1,358 @@ +// 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::fmt; +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, + } + } +} + +impl fmt::Display for BinaryFormat { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(match self { + Self::Hex => "hex", + Self::Utf8 => "utf-8", + Self::Base64 => "base64", + }) + } +} + +/// 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: Signature::one_of( + vec![TypeSignature::String(1), TypeSignature::String(2)], + Volatility::Immutable, + ), + } + } +} + +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: Signature::one_of( + vec![TypeSignature::String(1), TypeSignature::String(2)], + Volatility::Immutable, + ), + } + } +} + +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) + } +} + +/// 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 len = array.len(); + match array.data_type() { + DataType::Utf8 => convert_values( + name, + array.as_string::().iter(), + len, + format, + fail_on_error, + ), + DataType::LargeUtf8 => convert_values( + name, + array.as_string::().iter(), + len, + format, + fail_on_error, + ), + DataType::Utf8View => convert_values( + name, + array.as_string_view().iter(), + len, + format, + fail_on_error, + ), + other => exec_err!("{name}: expected a string argument, got {other}"), + } +} + +fn convert_values<'a>( + name: &str, + values: impl Iterator>, + len: usize, + format: BinaryFormat, + fail_on_error: bool, +) -> Result { + let mut builder = BinaryBuilder::with_capacity(len, 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 '{format}'" + ), + 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::*; + #[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..133bb19826672 100644 --- a/datafusion/sqllogictest/test_files/spark/string/to_binary.slt +++ b/datafusion/sqllogictest/test_files/spark/string/to_binary.slt @@ -23,5 +23,108 @@ ## 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 + +query ? +SELECT to_binary('abc'::string, 'utf8'::string); +---- +616263 + +# multi-byte character test +query ? +SELECT to_binary('é'::string, 'utf-8'::string); +---- +c3a9 + +# 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 + +# characters outside the base64 alphabet are skipped, as Java's MIME decoder does +query ? +SELECT to_binary('U3Bh cms='::string, 'base64'::string); +---- +537061726b + +# an unpadded final group is accepted +query ? +SELECT to_binary('U3Bhcms'::string, 'base64'::string); +---- +537061726b + +# empty input converts to empty binary in every format +query T +SELECT arrow_cast(to_binary(''::string, 'hex'::string), 'Utf8'); +---- +(empty) + +query T +SELECT arrow_cast(to_binary(''::string, 'utf-8'::string), 'Utf8'); +---- +(empty) + +query T +SELECT arrow_cast(to_binary(''::string, 'base64'::string), 'Utf8'); +---- +(empty) + +# 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 + +# a column of values is converted row by row, with NULLs preserved +query ? +SELECT to_binary(a, 'utf-8'::string) FROM (VALUES ('abc'::string), (NULL), ('déf')) AS t(a); +---- +616263 +NULL +64c3a966 + +# to_binary raises on a value that cannot be converted +statement error DataFusion error: Execution error: to_binary: cannot convert 'zz' to binary using format 'hex' +SELECT to_binary('zz'::string, 'hex'::string); + +statement error DataFusion error: Execution error: to_binary: cannot convert 'a!' to binary using format 'base64' +SELECT to_binary('a!'::string, 'base64'::string); + +# ...and on a format it does not recognise +statement error DataFusion error: Error during planning: to_binary: invalid fmt 'invalidFormat', expected one of 'hex', 'utf\-8', 'utf8' or 'base64' +SELECT to_binary('abc'::string, 'invalidFormat'::string); + +# one bad value fails the whole batch +statement error DataFusion error: Execution error: to_binary: cannot convert 'zz' to binary using format 'hex' +SELECT to_binary(a, 'hex'::string) FROM (VALUES ('4142'::string), ('zz')) AS t(a); 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..719e98d6483f4 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,50 @@ ## 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 + +# in a column, only the offending row becomes NULL +query ? +SELECT try_to_binary(a, 'hex'::string) FROM (VALUES ('4142'::string), ('zz')) AS t(a); +---- +4142 +NULL