diff --git a/datafusion/expr-common/src/signature.rs b/datafusion/expr-common/src/signature.rs index 2bde7c89a8259..d88cc791e99ca 100644 --- a/datafusion/expr-common/src/signature.rs +++ b/datafusion/expr-common/src/signature.rs @@ -508,6 +508,9 @@ impl TypeSignatureClass { TypeSignatureClass::Binary if native_type.is_binary() => { Ok(origin_type.to_owned()) } + // Binary has an unambiguous default, unlike Timestamp; cast NULL rather + // than letting the generic arm below pass DataType::Null downstream. + TypeSignatureClass::Binary if native_type.is_null() => Ok(DataType::Binary), TypeSignatureClass::Decimal if native_type.is_decimal() => { Ok(origin_type.to_owned()) } diff --git a/datafusion/functions/src/string/octet_length.rs b/datafusion/functions/src/string/octet_length.rs index 02df262ee27aa..90ea81cb70312 100644 --- a/datafusion/functions/src/string/octet_length.rs +++ b/datafusion/functions/src/string/octet_length.rs @@ -24,13 +24,13 @@ use datafusion_common::utils::take_function_args; use datafusion_common::{Result, ScalarValue}; use datafusion_expr::{ Coercion, ColumnarValue, Documentation, EncodingPreservation, ScalarFunctionArgs, - ScalarUDFImpl, Signature, TypeSignatureClass, Volatility, + ScalarUDFImpl, Signature, TypeSignature, TypeSignatureClass, Volatility, }; use datafusion_macros::user_doc; #[user_doc( doc_section(label = "String Functions"), - description = "Returns the length of a string in bytes.", + description = "Returns the length of a string or binary in bytes.", syntax_example = "octet_length(str)", sql_example = r#"```sql > select octet_length('Ångström'); @@ -40,7 +40,10 @@ use datafusion_macros::user_doc; | 10 | +--------------------------------+ ```"#, - standard_argument(name = "str", prefix = "String"), + argument( + name = "str", + description = "String or binary expression to operate on. Can be a constant, column, or function, and any combination of operators." + ), related_udf(name = "bit_length"), related_udf(name = "length") )] @@ -58,10 +61,22 @@ impl Default for OctetLengthFunc { impl OctetLengthFunc { pub fn new() -> Self { Self { - signature: Signature::coercible( + signature: Signature::one_of( vec![ - Coercion::new_exact(TypeSignatureClass::Native(logical_string())) - .with_encoding_preservation(EncodingPreservation::dictionary()), + TypeSignature::Coercible(vec![ + Coercion::new_exact(TypeSignatureClass::Native(logical_string())) + .with_encoding_preservation( + EncodingPreservation::dictionary(), + ), + ]), + // `TypeSignatureClass::Binary` also admits FixedSizeBinary, + // which `Native(logical_binary())` would reject. + TypeSignature::Coercible(vec![ + Coercion::new_exact(TypeSignatureClass::Binary) + .with_encoding_preservation( + EncodingPreservation::dictionary(), + ), + ]), ], Volatility::Immutable, ), @@ -107,6 +122,16 @@ fn octet_length_scalar(value: &ScalarValue) -> ScalarValue { ScalarValue::Utf8View(v) => { ScalarValue::Int32(v.as_ref().map(|x| x.len() as i32)) } + ScalarValue::Binary(v) => ScalarValue::Int32(v.as_ref().map(|x| x.len() as i32)), + ScalarValue::LargeBinary(v) => { + ScalarValue::Int64(v.as_ref().map(|x| x.len() as i64)) + } + ScalarValue::BinaryView(v) => { + ScalarValue::Int32(v.as_ref().map(|x| x.len() as i32)) + } + ScalarValue::FixedSizeBinary(_, v) => { + ScalarValue::Int32(v.as_ref().map(|x| x.len() as i32)) + } ScalarValue::Dictionary(key_type, value) => ScalarValue::Dictionary( key_type.clone(), Box::new(octet_length_scalar(value)), @@ -119,8 +144,11 @@ fn octet_length_scalar(value: &ScalarValue) -> ScalarValue { mod tests { use std::sync::Arc; - use arrow::array::{Array, Int32Array, StringArray}; - use arrow::datatypes::DataType::Int32; + use arrow::array::{ + Array, BinaryArray, BinaryViewArray, FixedSizeBinaryArray, Int32Array, + Int64Array, LargeBinaryArray, StringArray, + }; + use arrow::datatypes::DataType::{Int32, Int64}; use datafusion_common::ScalarValue; use datafusion_common::{Result, exec_err}; @@ -232,6 +260,130 @@ mod tests { Int32Array ); + // Binary inputs: byte length, no string coercion. + test_function!( + OctetLengthFunc::new(), + vec![ColumnarValue::Array(Arc::new(BinaryArray::from(vec![ + &b"chars"[..], + &b"chars2"[..], + ])))], + Ok(Some(5)), + i32, + Int32, + Int32Array + ); + test_function!( + OctetLengthFunc::new(), + vec![ColumnarValue::Scalar(ScalarValue::Binary(Some( + b"chars".to_vec() + )))], + Ok(Some(5)), + i32, + Int32, + Int32Array + ); + // Arbitrary non-UTF-8 bytes: the case CAST(col AS VARCHAR) cannot serve. + test_function!( + OctetLengthFunc::new(), + vec![ColumnarValue::Scalar(ScalarValue::Binary(Some(vec![ + 0xff, 0xfe, 0x00, 0x80 + ])))], + Ok(Some(4)), + i32, + Int32, + Int32Array + ); + test_function!( + OctetLengthFunc::new(), + vec![ColumnarValue::Scalar(ScalarValue::Binary(Some(vec![])))], + Ok(Some(0)), + i32, + Int32, + Int32Array + ); + test_function!( + OctetLengthFunc::new(), + vec![ColumnarValue::Scalar(ScalarValue::Binary(None))], + Ok(None), + i32, + Int32, + Int32Array + ); + test_function!( + OctetLengthFunc::new(), + vec![ColumnarValue::Array(Arc::new(BinaryViewArray::from(vec![ + &b"chars"[..], + &b"chars2"[..], + ])))], + Ok(Some(5)), + i32, + Int32, + Int32Array + ); + test_function!( + OctetLengthFunc::new(), + vec![ColumnarValue::Scalar(ScalarValue::BinaryView(Some( + b"chars".to_vec() + )))], + Ok(Some(5)), + i32, + Int32, + Int32Array + ); + // FixedSizeBinary reports its fixed width per non-null row. + test_function!( + OctetLengthFunc::new(), + vec![ColumnarValue::Array(Arc::new( + FixedSizeBinaryArray::try_from_iter( + vec![&b"abc"[..], &b"def"[..]].into_iter() + )? + ))], + Ok(Some(3)), + i32, + Int32, + Int32Array + ); + test_function!( + OctetLengthFunc::new(), + vec![ColumnarValue::Scalar(ScalarValue::FixedSizeBinary( + 5, + Some(b"chars".to_vec()) + ))], + Ok(Some(5)), + i32, + Int32, + Int32Array + ); + test_function!( + OctetLengthFunc::new(), + vec![ColumnarValue::Scalar(ScalarValue::FixedSizeBinary(5, None))], + Ok(None), + i32, + Int32, + Int32Array + ); + // LargeBinary widens the return type to Int64, mirroring LargeUtf8. + test_function!( + OctetLengthFunc::new(), + vec![ColumnarValue::Array(Arc::new(LargeBinaryArray::from( + vec![&b"chars"[..], &b"chars2"[..]] + )))], + Ok(Some(5)), + i64, + Int64, + Int64Array + ); + test_function!( + OctetLengthFunc::new(), + vec![ColumnarValue::Scalar(ScalarValue::LargeBinary(Some( + b"chars".to_vec() + )))], + Ok(Some(5)), + i64, + Int64, + Int64Array + ); + Ok(()) } } diff --git a/datafusion/functions/src/utils.rs b/datafusion/functions/src/utils.rs index b93bdb0b0d3bb..ad6fe32b868ba 100644 --- a/datafusion/functions/src/utils.rs +++ b/datafusion/functions/src/utils.rs @@ -32,7 +32,8 @@ use std::sync::Arc; /// /// If the input type is `Utf8` or `Binary` the return type is `$utf8Type`, /// -/// If the input type is `Utf8View` the return type is $utf8Type, +/// If the input type is `Utf8View`, `BinaryView` or `FixedSizeBinary` the +/// return type is `$utf8Type`, macro_rules! get_optimal_return_type { ($FUNC:ident, $largeUtf8Type:expr, $utf8Type:expr) => { pub(crate) fn $FUNC(arg_type: &DataType, name: &str) -> Result { @@ -43,10 +44,13 @@ macro_rules! get_optimal_return_type { DataType::Utf8 | DataType::Binary => $utf8Type, // Utf8View max offset size is u32::MAX, the same as UTF8 DataType::Utf8View | DataType::BinaryView => $utf8Type, + // FixedSizeBinary sizes are declared as i32 + DataType::FixedSizeBinary(_) => $utf8Type, DataType::Null => DataType::Null, DataType::Dictionary(_, value_type) => match **value_type { DataType::LargeUtf8 | DataType::LargeBinary => $largeUtf8Type, DataType::Utf8 | DataType::Binary => $utf8Type, + DataType::FixedSizeBinary(_) => $utf8Type, DataType::Null => DataType::Null, _ => { return datafusion_common::exec_err!( diff --git a/datafusion/sqllogictest/test_files/functions.slt b/datafusion/sqllogictest/test_files/functions.slt index 008be05852c85..486892d8adbc3 100644 --- a/datafusion/sqllogictest/test_files/functions.slt +++ b/datafusion/sqllogictest/test_files/functions.slt @@ -918,6 +918,71 @@ ORDER BY id 2 2 Dictionary(Int32, Int32) Dictionary(Int32, Dictionary(UInt32, Int32)) NULL NULL Dictionary(Int32, Int32) Dictionary(Int32, Dictionary(UInt32, Int32)) +# octet_length over binary types: byte semantics, no string coercion +query I +SELECT octet_length(arrow_cast('foo', 'Binary')) +---- +3 + +query IT +SELECT octet_length(arrow_cast('josé', 'Binary')), + arrow_typeof(octet_length(arrow_cast('josé', 'Binary'))) +---- +5 Int32 + +query IT +SELECT octet_length(arrow_cast('foo', 'BinaryView')), + arrow_typeof(octet_length(arrow_cast('foo', 'BinaryView'))) +---- +3 Int32 + +# LargeBinary widens the return type to Int64, mirroring LargeUtf8 +query IT +SELECT octet_length(arrow_cast('foo', 'LargeBinary')), + arrow_typeof(octet_length(arrow_cast('foo', 'LargeBinary'))) +---- +3 Int64 + +query I +SELECT octet_length(arrow_cast(NULL, 'Binary')) +---- +NULL + +# A bare NULL coerces through the string branch and keeps a typed Int32 result +query IT +SELECT octet_length(NULL), arrow_typeof(octet_length(NULL)) +---- +NULL Int32 + +# A dictionary over Null coerces to a string dictionary rather than reaching +# the arrow length kernel, which rejects Null values +query ?T +SELECT octet_length(arrow_cast(NULL, 'Dictionary(Int32, Null)')), + arrow_typeof(octet_length(arrow_cast(NULL, 'Dictionary(Int32, Null)'))) +---- +NULL Dictionary(Int32, Int32) + +# Dictionary encoding is preserved over binary values too +query ?T +SELECT octet_length(arrow_cast(arrow_cast('foo', 'Binary'), 'Dictionary(Int32, Binary)')), + arrow_typeof(octet_length(arrow_cast(arrow_cast('foo', 'Binary'), 'Dictionary(Int32, Binary)'))) +---- +3 Dictionary(Int32, Int32) + +# FixedSizeBinary reports its fixed width +query IT +SELECT octet_length(arrow_cast(arrow_cast('foo', 'Binary'), 'FixedSizeBinary(3)')), + arrow_typeof(octet_length(arrow_cast(arrow_cast('foo', 'Binary'), 'FixedSizeBinary(3)'))) +---- +3 Int32 + +# Dictionary encoding is preserved over FixedSizeBinary as well +query ?T +SELECT octet_length(arrow_cast(arrow_cast(arrow_cast('foo', 'Binary'), 'FixedSizeBinary(3)'), 'Dictionary(Int32, FixedSizeBinary(3))')), + arrow_typeof(octet_length(arrow_cast(arrow_cast(arrow_cast('foo', 'Binary'), 'FixedSizeBinary(3)'), 'Dictionary(Int32, FixedSizeBinary(3))'))) +---- +3 Dictionary(Int32, Int32) + query ??TT SELECT character_length(dict_col), character_length(nested_dict_col), arrow_typeof(character_length(dict_col)), diff --git a/docs/source/user-guide/expressions.md b/docs/source/user-guide/expressions.md index 3fbc11e0c92c0..55f64135dd0fb 100644 --- a/docs/source/user-guide/expressions.md +++ b/docs/source/user-guide/expressions.md @@ -191,7 +191,7 @@ select log(-1), log(0), sqrt(-1); | lpad(text, length, [, fill]) | Extends the string to length (`length`) by prepending the characters (`fill`) (a space by default). Example: `lpad('bb', 5, 'a') → aaabb` | | ltrim(text, text) | Removes all specified characters (`characters`) from the beginning of the string (`text`). Example: `ltrim('aabchelloccb', 'abc') -> helloccb` | | md5(text) | Computes the MD5 hash of the argument (`text`). | -| octet_length(text) | Returns number of bytes in the string (`text`). | +| octet_length(text) | Returns number of bytes in the string or binary (`text`). | | repeat(text, number) | Repeats the string the specified number of times. Example: `repeat('1', 4) -> 1111` | | replace(string, from, to) | Replaces a specified string (`from`) with another specified string (`to`) in the string (`string`). Example: `replace('Hello', 'replace', 'el') -> Hola` | | reverse(text) | Reverses the order of the characters in the string (`text`). Example: `reverse('hello') -> olleh` | diff --git a/docs/source/user-guide/sql/scalar_functions.md b/docs/source/user-guide/sql/scalar_functions.md index 644d42935ec29..cd21bb154809b 100644 --- a/docs/source/user-guide/sql/scalar_functions.md +++ b/docs/source/user-guide/sql/scalar_functions.md @@ -1663,7 +1663,7 @@ trim(LEADING trim_str FROM str) ### `octet_length` -Returns the length of a string in bytes. +Returns the length of a string or binary in bytes. ```sql octet_length(str) @@ -1671,7 +1671,7 @@ octet_length(str) #### Arguments -- **str**: String expression to operate on. Can be a constant, column, or function, and any combination of operators. +- **str**: String or binary expression to operate on. Can be a constant, column, or function, and any combination of operators. #### Example