diff --git a/datafusion/functions/Cargo.toml b/datafusion/functions/Cargo.toml index 83f8c0f2a3299..ebb80060daf6d 100644 --- a/datafusion/functions/Cargo.toml +++ b/datafusion/functions/Cargo.toml @@ -407,4 +407,4 @@ required-features = ["math_expressions"] [[bench]] harness = false name = "dictionary_encoding" -required-features = ["string_expressions"] +required-features = ["string_expressions", "unicode_expressions"] diff --git a/datafusion/functions/benches/dictionary_encoding.rs b/datafusion/functions/benches/dictionary_encoding.rs index 3afc1d5eb4c19..05541fc10e1d5 100644 --- a/datafusion/functions/benches/dictionary_encoding.rs +++ b/datafusion/functions/benches/dictionary_encoding.rs @@ -42,10 +42,16 @@ fn create_string_dictionary(cardinality: usize) -> ArrayRef { } fn benchmark_dictionary_string_udfs(c: &mut Criterion) { - let udfs: [(&str, Arc); 3] = [ + let udfs: [(&str, Arc); 6] = [ ("ascii", datafusion_functions::string::ascii()), ("bit_length", datafusion_functions::string::bit_length()), + ( + "character_length", + datafusion_functions::unicode::character_length(), + ), + ("initcap", datafusion_functions::unicode::initcap()), ("octet_length", datafusion_functions::string::octet_length()), + ("reverse", datafusion_functions::unicode::reverse()), ]; let config_options = Arc::new(ConfigOptions::default()); diff --git a/datafusion/functions/src/unicode/character_length.rs b/datafusion/functions/src/unicode/character_length.rs index 465b15ace1d10..aa38a0a1d2703 100644 --- a/datafusion/functions/src/unicode/character_length.rs +++ b/datafusion/functions/src/unicode/character_length.rs @@ -15,16 +15,19 @@ // specific language governing permissions and limitations // under the License. -use crate::utils::{make_scalar_function, utf8_to_int_type}; +use crate::utils::{ + make_scalar_function, transform_leaf_type_preserving_encoding, utf8_to_int_type, +}; use arrow::array::{ Array, ArrayRef, ArrowPrimitiveType, AsArray, OffsetSizeTrait, PrimitiveArray, StringArrayType, }; use arrow::datatypes::{ArrowNativeType, DataType, Int32Type, Int64Type}; use datafusion_common::Result; +use datafusion_common::types::{NativeType, logical_binary, logical_string}; use datafusion_expr::{ - ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, - Volatility, + Coercion, ColumnarValue, Documentation, EncodingPreservation, ScalarFunctionArgs, + ScalarUDFImpl, Signature, TypeSignatureClass, Volatility, }; use datafusion_macros::user_doc; use std::sync::Arc; @@ -59,11 +62,16 @@ impl Default for CharacterLengthFunc { impl CharacterLengthFunc { pub fn new() -> Self { - use DataType::*; Self { - signature: Signature::uniform( - 1, - vec![Utf8, LargeUtf8, Utf8View], + signature: Signature::coercible( + vec![ + Coercion::new_implicit( + TypeSignatureClass::Native(logical_string()), + vec![TypeSignatureClass::Native(logical_binary())], + NativeType::String, + ) + .with_encoding_preservation(EncodingPreservation::dictionary()), + ], Volatility::Immutable, ), aliases: vec![String::from("length"), String::from("char_length")], @@ -81,7 +89,9 @@ impl ScalarUDFImpl for CharacterLengthFunc { } fn return_type(&self, arg_types: &[DataType]) -> Result { - utf8_to_int_type(&arg_types[0], "character_length") + transform_leaf_type_preserving_encoding(&arg_types[0], &|data_type| { + utf8_to_int_type(data_type, "character_length") + }) } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { @@ -114,6 +124,11 @@ fn character_length(args: &[ArrayRef]) -> Result { let string_array = args[0].as_string_view(); character_length_general::(&string_array) } + DataType::Dictionary(_, _) => { + let dictionary = args[0].as_any_dictionary(); + let converted = character_length(&[Arc::clone(dictionary.values())])?; + Ok(dictionary.with_values(converted)) + } _ => unreachable!("CharacterLengthFunc"), } } diff --git a/datafusion/functions/src/unicode/initcap.rs b/datafusion/functions/src/unicode/initcap.rs index 8981d59aec8d2..0332ab5d4427f 100644 --- a/datafusion/functions/src/unicode/initcap.rs +++ b/datafusion/functions/src/unicode/initcap.rs @@ -17,18 +17,17 @@ use std::sync::Arc; -use arrow::array::{Array, ArrayRef, GenericStringArray, OffsetSizeTrait}; +use arrow::array::{Array, ArrayRef, AsArray, GenericStringArray, OffsetSizeTrait}; use arrow::buffer::Buffer; use arrow::datatypes::DataType; use crate::strings::{GenericStringArrayBuilder, StringViewArrayBuilder}; -use crate::utils::{make_scalar_function, utf8_to_str_type}; use datafusion_common::cast::{as_generic_string_array, as_string_view_array}; use datafusion_common::types::logical_string; use datafusion_common::{Result, ScalarValue, exec_err}; use datafusion_expr::{ - Coercion, ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, - TypeSignatureClass, Volatility, + Coercion, ColumnarValue, Documentation, EncodingPreservation, ScalarFunctionArgs, + ScalarUDFImpl, Signature, TypeSignatureClass, Volatility, }; use datafusion_macros::user_doc; @@ -64,9 +63,10 @@ impl InitcapFunc { pub fn new() -> Self { Self { signature: Signature::coercible( - vec![Coercion::new_exact(TypeSignatureClass::Native( - logical_string(), - ))], + vec![ + Coercion::new_exact(TypeSignatureClass::Native(logical_string())) + .with_encoding_preservation(EncodingPreservation::dictionary()), + ], Volatility::Immutable, ), } @@ -83,54 +83,16 @@ impl ScalarUDFImpl for InitcapFunc { } fn return_type(&self, arg_types: &[DataType]) -> Result { - if let DataType::Utf8View = arg_types[0] { - Ok(DataType::Utf8View) - } else { - utf8_to_str_type(&arg_types[0], "initcap") - } + Ok(arg_types[0].clone()) } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - let arg = &args.args[0]; - - // Scalar fast path - handle directly without array conversion - if let ColumnarValue::Scalar(scalar) = arg { - return match scalar { - ScalarValue::Utf8(None) - | ScalarValue::LargeUtf8(None) - | ScalarValue::Utf8View(None) => Ok(arg.clone()), - ScalarValue::Utf8(Some(s)) => { - let mut result = String::new(); - initcap_string(s, &mut result); - Ok(ColumnarValue::Scalar(ScalarValue::Utf8(Some(result)))) - } - ScalarValue::LargeUtf8(Some(s)) => { - let mut result = String::new(); - initcap_string(s, &mut result); - Ok(ColumnarValue::Scalar(ScalarValue::LargeUtf8(Some(result)))) - } - ScalarValue::Utf8View(Some(s)) => { - let mut result = String::new(); - initcap_string(s, &mut result); - Ok(ColumnarValue::Scalar(ScalarValue::Utf8View(Some(result)))) - } - other => { - exec_err!( - "Unsupported data type {:?} for function `initcap`", - other.data_type() - ) - } - }; - } - - // Array path - let args = &args.args; - match args[0].data_type() { - DataType::Utf8 => make_scalar_function(initcap::, vec![])(args), - DataType::LargeUtf8 => make_scalar_function(initcap::, vec![])(args), - DataType::Utf8View => make_scalar_function(initcap_utf8view, vec![])(args), - other => { - exec_err!("Unsupported data type {other:?} for function `initcap`") + match &args.args[0] { + ColumnarValue::Scalar(scalar) => { + Ok(ColumnarValue::Scalar(initcap_scalar(scalar)?)) + } + ColumnarValue::Array(array) => { + Ok(ColumnarValue::Array(initcap_array(array)?)) } } } @@ -140,6 +102,55 @@ impl ScalarUDFImpl for InitcapFunc { } } +fn initcap_scalar(scalar: &ScalarValue) -> Result { + match scalar { + ScalarValue::Utf8(None) + | ScalarValue::LargeUtf8(None) + | ScalarValue::Utf8View(None) => Ok(scalar.clone()), + ScalarValue::Utf8(Some(s)) => { + let mut result = String::new(); + initcap_string(s, &mut result); + Ok(ScalarValue::Utf8(Some(result))) + } + ScalarValue::LargeUtf8(Some(s)) => { + let mut result = String::new(); + initcap_string(s, &mut result); + Ok(ScalarValue::LargeUtf8(Some(result))) + } + ScalarValue::Utf8View(Some(s)) => { + let mut result = String::new(); + initcap_string(s, &mut result); + Ok(ScalarValue::Utf8View(Some(result))) + } + ScalarValue::Dictionary(key_type, value) => Ok(ScalarValue::Dictionary( + key_type.clone(), + Box::new(initcap_scalar(value)?), + )), + other => { + exec_err!( + "Unsupported data type {:?} for function `initcap`", + other.data_type() + ) + } + } +} + +fn initcap_array(array: &ArrayRef) -> Result { + match array.data_type() { + DataType::Utf8 => initcap::(&[Arc::clone(array)]), + DataType::LargeUtf8 => initcap::(&[Arc::clone(array)]), + DataType::Utf8View => initcap_utf8view(&[Arc::clone(array)]), + DataType::Dictionary(_, _) => { + let dictionary = array.as_any_dictionary(); + let converted = initcap_array(dictionary.values())?; + Ok(dictionary.with_values(converted)) + } + other => { + exec_err!("Unsupported data type {other:?} for function `initcap`") + } + } +} + /// Converts the first letter of each word to uppercase and the rest to /// lowercase. Words are sequences of alphanumeric characters separated by /// non-alphanumeric characters. diff --git a/datafusion/functions/src/unicode/reverse.rs b/datafusion/functions/src/unicode/reverse.rs index 813dcb5f504dd..1d802a3894b25 100644 --- a/datafusion/functions/src/unicode/reverse.rs +++ b/datafusion/functions/src/unicode/reverse.rs @@ -15,6 +15,8 @@ // specific language governing permissions and limitations // under the License. +use std::sync::Arc; + use crate::strings::{ BulkNullStringArrayBuilder, GenericStringArrayBuilder, StringViewArrayBuilder, }; @@ -22,10 +24,11 @@ use crate::utils::make_scalar_function; use DataType::{LargeUtf8, Utf8, Utf8View}; use arrow::array::{Array, ArrayRef, AsArray, StringArrayType}; use arrow::datatypes::DataType; +use datafusion_common::types::{NativeType, logical_binary, logical_string}; use datafusion_common::{Result, exec_err}; use datafusion_expr::{ - ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, - Volatility, + Coercion, ColumnarValue, Documentation, EncodingPreservation, ScalarFunctionArgs, + ScalarUDFImpl, Signature, TypeSignatureClass, Volatility, }; use datafusion_macros::user_doc; @@ -56,11 +59,16 @@ impl Default for ReverseFunc { impl ReverseFunc { pub fn new() -> Self { - use DataType::*; Self { - signature: Signature::uniform( - 1, - vec![Utf8View, Utf8, LargeUtf8], + signature: Signature::coercible( + vec![ + Coercion::new_implicit( + TypeSignatureClass::Native(logical_string()), + vec![TypeSignatureClass::Native(logical_binary())], + NativeType::String, + ) + .with_encoding_preservation(EncodingPreservation::dictionary()), + ], Volatility::Immutable, ), } @@ -83,7 +91,9 @@ impl ScalarUDFImpl for ReverseFunc { fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { let args = &args.args; match args[0].data_type() { - Utf8 | Utf8View | LargeUtf8 => make_scalar_function(reverse, vec![])(args), + Utf8 | Utf8View | LargeUtf8 | DataType::Dictionary(_, _) => { + make_scalar_function(reverse, vec![])(args) + } other => { exec_err!("Unsupported data type {other:?} for function reverse") } @@ -113,6 +123,11 @@ fn reverse(args: &[ArrayRef]) -> Result { &args[0].as_string_view(), StringViewArrayBuilder::with_capacity(len), ), + DataType::Dictionary(_, _) => { + let dictionary = args[0].as_any_dictionary(); + let converted = reverse(&[Arc::clone(dictionary.values())])?; + Ok(dictionary.with_values(converted)) + } _ => unreachable!( "Reverse can only be applied to Utf8View, Utf8 and LargeUtf8 types" ), diff --git a/datafusion/sqllogictest/test_files/binary.slt b/datafusion/sqllogictest/test_files/binary.slt index 94c1365cb9514..91a9449343d2a 100644 --- a/datafusion/sqllogictest/test_files/binary.slt +++ b/datafusion/sqllogictest/test_files/binary.slt @@ -281,7 +281,7 @@ SELECT cast(binary as varchar) as str, character_length(binary) as binary_len, cast(largebinary as varchar) as large_str, - character_length(binary) as largebinary_len + character_length(largebinary) as largebinary_len from t; ---- Foo 3 Foo 3 @@ -298,6 +298,20 @@ SELECT character_length(X'20'); query error Encountered non UTF\-8 data: invalid utf\-8 sequence of 1 bytes from index 0 SELECT character_length(X'c328'); +# reverse function +query TTTT +SELECT + cast(binary as varchar) as str, + reverse(binary) as binary_reversed, + cast(largebinary as varchar) as large_str, + reverse(largebinary) as largebinary_reversed +from t; +---- +Foo ooF Foo ooF +NULL NULL NULL NULL +Bar raB Bar raB +FooBar raBooF FooBar raBooF + # regexp_replace query TTTT SELECT @@ -363,4 +377,4 @@ hellohello query T SELECT 'hello' || arrow_cast(arrow_cast('hello', 'Binary'), 'BinaryView'); ---- -hellohello \ No newline at end of file +hellohello diff --git a/datafusion/sqllogictest/test_files/functions.slt b/datafusion/sqllogictest/test_files/functions.slt index 78045936a1893..78bdeb3e15520 100644 --- a/datafusion/sqllogictest/test_files/functions.slt +++ b/datafusion/sqllogictest/test_files/functions.slt @@ -68,7 +68,7 @@ SELECT length('') ---- 0 -query I +query ? SELECT length(arrow_cast('', 'Dictionary(Int32, Utf8)')) ---- 0 @@ -83,7 +83,7 @@ SELECT length('josé') ---- 4 -query I +query ? SELECT length(arrow_cast('josé', 'Dictionary(Int32, Utf8)')) ---- 4 @@ -507,6 +507,84 @@ SELECT initcap(arrow_cast('foo', 'Dictionary(Int32, Utf8)')) ---- Foo +query TTTT +SELECT initcap(arrow_cast('foo BAR', 'Dictionary(Int32, LargeUtf8)')), + arrow_typeof(initcap(arrow_cast('foo BAR', 'Dictionary(Int32, LargeUtf8)'))), + initcap(arrow_cast( + arrow_cast('foo BAR', 'Dictionary(Int32, Utf8)'), + 'Dictionary(Int32, Utf8View)' + )), + arrow_typeof(initcap(arrow_cast( + arrow_cast('foo BAR', 'Dictionary(Int32, Utf8)'), + 'Dictionary(Int32, Utf8View)' + ))) +---- +Foo Bar Dictionary(Int32, LargeUtf8) Foo Bar Dictionary(Int32, Utf8View) + +query ?T +SELECT initcap(arrow_cast( + arrow_cast('foo BAR', 'Dictionary(UInt32, Utf8)'), + 'Dictionary(Int32, Dictionary(UInt32, Utf8))' + )), + arrow_typeof(initcap(arrow_cast( + arrow_cast('foo BAR', 'Dictionary(UInt32, Utf8)'), + 'Dictionary(Int32, Dictionary(UInt32, Utf8))' + ))) +---- +Foo Bar Dictionary(Int32, Dictionary(UInt32, Utf8)) + +statement ok +CREATE TABLE unicode_dictionary_test AS +SELECT column1 AS id, + arrow_cast(column2, 'Dictionary(Int32, Utf8)') AS dict_col, + arrow_cast( + arrow_cast(column2, 'Dictionary(UInt32, Utf8)'), + 'Dictionary(Int32, Dictionary(UInt32, Utf8))' + ) AS nested_dict_col +FROM (VALUES +(1, 'foo BAR'), +(2, 'éclair CAFÉ'), +(3, NULL)); + +query T?TT +SELECT initcap(dict_col), initcap(nested_dict_col), + arrow_typeof(initcap(dict_col)), + arrow_typeof(initcap(nested_dict_col)) +FROM unicode_dictionary_test +ORDER BY id +---- +Foo Bar Foo Bar Dictionary(Int32, Utf8) Dictionary(Int32, Dictionary(UInt32, Utf8)) +Éclair Café Éclair Café Dictionary(Int32, Utf8) Dictionary(Int32, Dictionary(UInt32, Utf8)) +NULL NULL Dictionary(Int32, Utf8) Dictionary(Int32, Dictionary(UInt32, Utf8)) + +query TTTT +SELECT reverse(arrow_cast('foo BAR', 'Dictionary(Int32, LargeUtf8)')), + arrow_typeof(reverse(arrow_cast('foo BAR', 'Dictionary(Int32, LargeUtf8)'))), + reverse(arrow_cast( + arrow_cast('foo BAR', 'Dictionary(Int32, Utf8)'), + 'Dictionary(Int32, Utf8View)' + )), + arrow_typeof(reverse(arrow_cast( + arrow_cast('foo BAR', 'Dictionary(Int32, Utf8)'), + 'Dictionary(Int32, Utf8View)' + ))) +---- +RAB oof Dictionary(Int32, LargeUtf8) RAB oof Dictionary(Int32, Utf8View) + +query T?TT +SELECT reverse(dict_col), reverse(nested_dict_col), + arrow_typeof(reverse(dict_col)), + arrow_typeof(reverse(nested_dict_col)) +FROM unicode_dictionary_test +ORDER BY id +---- +RAB oof RAB oof Dictionary(Int32, Utf8) Dictionary(Int32, Dictionary(UInt32, Utf8)) +ÉFAC rialcé ÉFAC rialcé Dictionary(Int32, Utf8) Dictionary(Int32, Dictionary(UInt32, Utf8)) +NULL NULL Dictionary(Int32, Utf8) Dictionary(Int32, Dictionary(UInt32, Utf8)) + +statement ok +DROP TABLE unicode_dictionary_test + query ? SELECT ascii(arrow_cast('é', 'Dictionary(Int32, Utf8)')) ---- @@ -688,11 +766,39 @@ SELECT character_length('foo') ---- 3 -query I +query ? SELECT character_length(arrow_cast('foo', 'Dictionary(Int32, Utf8)')) ---- 3 +query ?T +SELECT character_length(arrow_cast( + arrow_cast('é', 'Dictionary(UInt32, Utf8)'), + 'Dictionary(Int32, Dictionary(UInt32, Utf8))' + )), + arrow_typeof(character_length(arrow_cast( + arrow_cast('é', 'Dictionary(UInt32, Utf8)'), + 'Dictionary(Int32, Dictionary(UInt32, Utf8))' + ))) +---- +1 Dictionary(Int32, Dictionary(UInt32, Int32)) + +query ?T?T +SELECT character_length(arrow_cast('foo', 'Dictionary(Int32, LargeUtf8)')), + arrow_typeof(character_length( + arrow_cast('foo', 'Dictionary(Int32, LargeUtf8)') + )), + character_length(arrow_cast( + arrow_cast('foo', 'Dictionary(Int32, Utf8)'), + 'Dictionary(Int32, Utf8View)' + )), + arrow_typeof(character_length(arrow_cast( + arrow_cast('foo', 'Dictionary(Int32, Utf8)'), + 'Dictionary(Int32, Utf8View)' + ))) +---- +3 Dictionary(Int32, Int64) 3 Dictionary(Int32, Int32) + query I SELECT octet_length('foo') ---- @@ -755,6 +861,17 @@ ORDER BY id 2 2 Dictionary(Int32, Int32) Dictionary(Int32, Dictionary(UInt32, Int32)) NULL NULL Dictionary(Int32, Int32) Dictionary(Int32, Dictionary(UInt32, Int32)) +query ??TT +SELECT character_length(dict_col), character_length(nested_dict_col), + arrow_typeof(character_length(dict_col)), + arrow_typeof(character_length(nested_dict_col)) +FROM string_length_dictionary_test +ORDER BY id +---- +3 3 Dictionary(Int32, Int32) Dictionary(Int32, Dictionary(UInt32, Int32)) +1 1 Dictionary(Int32, Int32) Dictionary(Int32, Dictionary(UInt32, Int32)) +NULL NULL Dictionary(Int32, Int32) Dictionary(Int32, Dictionary(UInt32, Int32)) + query ??TT SELECT ascii(dict_col), ascii(nested_dict_col), arrow_typeof(ascii(dict_col)), diff --git a/datafusion/sqllogictest/test_files/string/string_literal.slt b/datafusion/sqllogictest/test_files/string/string_literal.slt index c175f52a35f99..2b3e54aa9bdcf 100644 --- a/datafusion/sqllogictest/test_files/string/string_literal.slt +++ b/datafusion/sqllogictest/test_files/string/string_literal.slt @@ -486,6 +486,11 @@ SELECT reverse(arrow_cast('abcde', 'Dictionary(Int32, Utf8)')) ---- edcba +query T +SELECT arrow_typeof(reverse(arrow_cast('abcde', 'Dictionary(Int32, Utf8)'))) +---- +Dictionary(Int32, Utf8) + query T SELECT reverse('loẅks') ---- diff --git a/datafusion/sqllogictest/test_files/string/string_query.slt.part b/datafusion/sqllogictest/test_files/string/string_query.slt.part index 9231ec7b9c976..199d8a5f4d06c 100644 --- a/datafusion/sqllogictest/test_files/string/string_query.slt.part +++ b/datafusion/sqllogictest/test_files/string/string_query.slt.part @@ -1253,8 +1253,8 @@ NULL NULL NULL NULL NULL NULL query II SELECT - CHARACTER_LENGTH(ascii_1), - CHARACTER_LENGTH(unicode_1) + arrow_cast(CHARACTER_LENGTH(ascii_1), 'Int64'), + arrow_cast(CHARACTER_LENGTH(unicode_1), 'Int64') FROM test_basic_operator ----