diff --git a/datafusion/expr-common/src/signature.rs b/datafusion/expr-common/src/signature.rs index 35a679cc447cf..f0010f0a05014 100644 --- a/datafusion/expr-common/src/signature.rs +++ b/datafusion/expr-common/src/signature.rs @@ -1049,6 +1049,8 @@ pub enum Coercion { Exact { /// The required type for the argument desired_type: TypeSignatureClass, + /// Physical encoding preservation requested by the function. + encoding_preservation: EncodingPreservation, }, /// Coercion that accepts the desired type and can implicitly coerce from other types. @@ -1057,12 +1059,44 @@ pub enum Coercion { desired_type: TypeSignatureClass, /// Rules for implicit coercion from other types implicit_coercion: ImplicitCoercion, + /// Physical encoding preservation requested by the function. + encoding_preservation: EncodingPreservation, }, } +/// Controls whether a [`Coercion`] preserves an argument's physical encoding +/// (e.g. dictionary) instead of materializing it to the coerced value type. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, PartialOrd, Hash)] +pub struct EncodingPreservation { + preserve_dictionary: bool, +} + +impl EncodingPreservation { + /// Preserve dictionary encoding and coerce only the dictionary values. + pub const fn dictionary() -> Self { + Self { + preserve_dictionary: true, + } + } + + /// Preserve dictionary encoding and coerce only the dictionary values. + pub const fn with_dictionary(mut self) -> Self { + self.preserve_dictionary = true; + self + } + + /// Returns whether dictionary encoding should be preserved. + pub const fn preserve_dictionary(self) -> bool { + self.preserve_dictionary + } +} + impl Coercion { pub fn new_exact(desired_type: TypeSignatureClass) -> Self { - Self::Exact { desired_type } + Self::Exact { + desired_type, + encoding_preservation: EncodingPreservation::default(), + } } /// Create a new coercion with implicit coercion rules. @@ -1080,6 +1114,37 @@ impl Coercion { allowed_source_types, default_casted_type, }, + encoding_preservation: EncodingPreservation::default(), + } + } + + pub fn with_encoding_preservation( + mut self, + encoding_preservation: EncodingPreservation, + ) -> Self { + match &mut self { + Coercion::Exact { + encoding_preservation: current, + .. + } + | Coercion::Implicit { + encoding_preservation: current, + .. + } => *current = encoding_preservation, + } + self + } + + pub fn encoding_preservation(&self) -> EncodingPreservation { + match self { + Coercion::Exact { + encoding_preservation, + .. + } + | Coercion::Implicit { + encoding_preservation, + .. + } => *encoding_preservation, } } @@ -1103,7 +1168,7 @@ impl Coercion { pub fn desired_type(&self) -> &TypeSignatureClass { match self { - Coercion::Exact { desired_type } => desired_type, + Coercion::Exact { desired_type, .. } => desired_type, Coercion::Implicit { desired_type, .. } => desired_type, } } @@ -1128,6 +1193,7 @@ impl PartialEq for Coercion { fn eq(&self, other: &Self) -> bool { self.desired_type() == other.desired_type() && self.implicit_coercion() == other.implicit_coercion() + && self.encoding_preservation() == other.encoding_preservation() } } @@ -1135,6 +1201,7 @@ impl Hash for Coercion { fn hash(&self, state: &mut H) { self.desired_type().hash(state); self.implicit_coercion().hash(state); + self.encoding_preservation().hash(state); } } @@ -2178,6 +2245,20 @@ mod tests { assert_snapshot!(implicit_with_multiple_sources, @"Int64"); } + #[test] + fn test_coercion_encoding_preservation_affects_equality() { + assert!(!EncodingPreservation::default().preserve_dictionary()); + let preserve_dictionary = EncodingPreservation::dictionary(); + assert!(preserve_dictionary.preserve_dictionary()); + + let default = Coercion::new_exact(TypeSignatureClass::Native(logical_string())); + let preserving = default + .clone() + .with_encoding_preservation(preserve_dictionary); + + assert_ne!(default, preserving); + } + #[test] fn test_to_string_repr_coercible() { use insta::assert_snapshot; diff --git a/datafusion/expr/src/lib.rs b/datafusion/expr/src/lib.rs index b52a784df931a..43cb3fdc20c40 100644 --- a/datafusion/expr/src/lib.rs +++ b/datafusion/expr/src/lib.rs @@ -101,8 +101,8 @@ pub use datafusion_expr_common::groups_accumulator::{EmitTo, GroupsAccumulator}; pub use datafusion_expr_common::operator::Operator; pub use datafusion_expr_common::placement::ExpressionPlacement; pub use datafusion_expr_common::signature::{ - ArrayFunctionArgument, ArrayFunctionSignature, Coercion, Signature, - TIMEZONE_WILDCARD, TypeSignature, TypeSignatureClass, Volatility, + ArrayFunctionArgument, ArrayFunctionSignature, Coercion, EncodingPreservation, + Signature, TIMEZONE_WILDCARD, TypeSignature, TypeSignatureClass, Volatility, }; pub use datafusion_expr_common::type_coercion::binary; pub use expr::{ diff --git a/datafusion/expr/src/type_coercion/functions.rs b/datafusion/expr/src/type_coercion/functions.rs index 33746a2c46b30..c2dc56ae1008a 100644 --- a/datafusion/expr/src/type_coercion/functions.rs +++ b/datafusion/expr/src/type_coercion/functions.rs @@ -33,7 +33,7 @@ use datafusion_common::utils::{ use datafusion_common::{ Result, exec_err, internal_err, plan_err, types::NativeType, utils::list_ndims, }; -use datafusion_expr_common::signature::ArrayFunctionArgument; +use datafusion_expr_common::signature::{ArrayFunctionArgument, EncodingPreservation}; use datafusion_expr_common::type_coercion::binary::type_union_resolution; use datafusion_expr_common::{ signature::{ArrayFunctionSignature, FIXED_SIZE_LIST_WILDCARD, TIMEZONE_WILDCARD}, @@ -873,9 +873,39 @@ fn get_valid_types( TypeSignature::Coercible(param_types) => { function_length_check(function_name, current_types.len(), param_types.len())?; + fn cast_origin( + current_type: &DataType, + encoding_preservation: EncodingPreservation, + ) -> &DataType { + if encoding_preservation.preserve_dictionary() + && let DataType::Dictionary(_, value_type) = current_type + { + value_type + } else { + current_type + } + } + + fn preserve_encoding( + current_type: &DataType, + casted_type: DataType, + encoding_preservation: EncodingPreservation, + ) -> DataType { + if encoding_preservation.preserve_dictionary() + && let DataType::Dictionary(key_type, _) = current_type + && !matches!(casted_type, DataType::Dictionary(_, _)) + { + DataType::Dictionary(key_type.clone(), Box::new(casted_type)) + } else { + casted_type + } + } + let mut new_types = Vec::with_capacity(current_types.len()); for (current_type, param) in current_types.iter().zip(param_types.iter()) { let current_native_type: NativeType = current_type.into(); + let encoding_preservation = param.encoding_preservation(); + let cast_origin = cast_origin(current_type, encoding_preservation); if param .desired_type() @@ -883,9 +913,13 @@ fn get_valid_types( { let casted_type = param .desired_type() - .default_casted_type(¤t_native_type, current_type)?; + .default_casted_type(¤t_native_type, cast_origin)?; - new_types.push(casted_type); + new_types.push(preserve_encoding( + current_type, + casted_type, + encoding_preservation, + )); } else if param .allowed_source_types() .iter() @@ -894,8 +928,12 @@ fn get_valid_types( // If the condition is met which means `implicit coercion`` is provided so we can safely unwrap let default_casted_type = param.default_casted_type().unwrap(); let casted_type = - default_casted_type.default_cast_for(current_type)?; - new_types.push(casted_type); + default_casted_type.default_cast_for(cast_origin)?; + new_types.push(preserve_encoding( + current_type, + casted_type, + encoding_preservation, + )); } else { let hint = if matches!(current_native_type, NativeType::Binary) { "\n\nHint: Binary types are not automatically coerced to String. Use CAST(column AS VARCHAR) to convert Binary data to String." @@ -1214,11 +1252,11 @@ mod tests { use arrow::datatypes::IntervalUnit; use datafusion_common::{ assert_contains, - types::{logical_binary, logical_int64}, + types::{logical_binary, logical_int64, logical_string}, }; use datafusion_expr_common::{ columnar_value::ColumnarValue, - signature::{Coercion, TypeSignatureClass}, + signature::{Coercion, EncodingPreservation, TypeSignatureClass}, }; #[test] @@ -1831,6 +1869,112 @@ mod tests { Ok(()) } + #[test] + fn test_coercible_dictionary_preserves_encoding() -> Result<()> { + fn dictionary_input( + value_type: DataType, + coercion: Coercion, + ) -> Result> { + fields_with_udf( + &[Field::new( + "field", + DataType::Dictionary(Box::new(DataType::Int8), Box::new(value_type)), + true, + ) + .into()], + &MockUdf(Signature::coercible(vec![coercion], Volatility::Immutable)), + ) + .map(|v| v.into_iter().map(|f| f.data_type().clone()).collect()) + } + + let coercion = Coercion::new_exact(TypeSignatureClass::Native(logical_string())) + .with_encoding_preservation(EncodingPreservation::dictionary()); + + assert_eq!( + dictionary_input(DataType::LargeUtf8, coercion.clone())?, + vec![DataType::Dictionary( + Box::new(DataType::Int8), + Box::new(DataType::LargeUtf8), + )] + ); + assert_eq!( + dictionary_input( + DataType::BinaryView, + Coercion::new_implicit( + TypeSignatureClass::Native(logical_string()), + vec![TypeSignatureClass::Native(logical_binary())], + NativeType::String, + ) + .with_encoding_preservation(EncodingPreservation::dictionary()), + )?, + vec![DataType::Dictionary( + Box::new(DataType::Int8), + Box::new(DataType::Utf8View), + )] + ); + // Contrast: without encoding_preservation, Native strips dictionary entirely + assert_eq!( + dictionary_input( + DataType::Int32, + Coercion::new_implicit( + TypeSignatureClass::Native(logical_int64()), + vec![TypeSignatureClass::Integer], + NativeType::Int64, + ), + )?, + vec![DataType::Int64] + ); + // With encoding_preservation, dictionary wrapper is preserved, value coerced + assert_eq!( + dictionary_input( + DataType::Int32, + Coercion::new_implicit( + TypeSignatureClass::Native(logical_int64()), + vec![TypeSignatureClass::Integer], + NativeType::Int64, + ) + .with_encoding_preservation(EncodingPreservation::dictionary()), + )?, + vec![DataType::Dictionary( + Box::new(DataType::Int8), + Box::new(DataType::Int64), + )] + ); + // Contrast: without encoding_preservation, non-Native already passes through + assert_eq!( + dictionary_input( + DataType::Int32, + Coercion::new_implicit( + TypeSignatureClass::Integer, + vec![], + NativeType::Int64, + ), + )?, + vec![DataType::Dictionary( + Box::new(DataType::Int8), + Box::new(DataType::Int32), + )] + ); + // With encoding_preservation, same result — no difference for non-Native + assert_eq!( + dictionary_input( + DataType::Int32, + Coercion::new_implicit( + TypeSignatureClass::Integer, + vec![], + NativeType::Int64, + ) + .with_encoding_preservation(EncodingPreservation::dictionary()), + )?, + vec![DataType::Dictionary( + Box::new(DataType::Int8), + Box::new(DataType::Int32), + )] + ); + + Ok(()) + } + #[test] fn test_coercible_run_end_encoded() -> Result<()> { let run_end_encoded = DataType::RunEndEncoded( diff --git a/datafusion/functions/src/string/common.rs b/datafusion/functions/src/string/common.rs index 729c151f4dc3e..b51b92e9df1ed 100644 --- a/datafusion/functions/src/string/common.rs +++ b/datafusion/functions/src/string/common.rs @@ -24,7 +24,7 @@ use crate::strings::{ StringViewArrayBuilder, append_view, }; use arrow::array::{ - Array, ArrayRef, GenericStringArray, NullBufferBuilder, OffsetSizeTrait, + Array, ArrayRef, AsArray, GenericStringArray, NullBufferBuilder, OffsetSizeTrait, StringViewArray, new_null_array, }; use arrow::buffer::{Buffer, OffsetBuffer, ScalarBuffer}; @@ -348,63 +348,100 @@ fn case_conversion( name: &str, ) -> Result { match &args[0] { - ColumnarValue::Array(array) => match array.data_type() { - DataType::Utf8 => Ok(ColumnarValue::Array(case_conversion_array::( - array, lower, - )?)), - DataType::LargeUtf8 => Ok(ColumnarValue::Array( - case_conversion_array::(array, lower)?, - )), - DataType::Utf8View => { - let string_array = as_string_view_array(array)?; - if string_array.is_ascii() { - return Ok(ColumnarValue::Array(Arc::new( - case_conversion_utf8view_ascii(string_array, lower), - ))); - } - let item_len = string_array.len(); - // Null-preserving: reuse the input null buffer as the output null buffer. - let nulls = string_array.nulls().cloned(); - let mut builder = StringViewArrayBuilder::with_capacity(item_len); - - if let Some(ref n) = nulls { - for i in 0..item_len { - if n.is_null(i) { - builder.try_append_placeholder()?; - } else { - // SAFETY: `n.is_null(i)` was false in the branch above. - let s = unsafe { string_array.value_unchecked(i) }; - builder.try_append_value(&unicode_case(s, lower))?; - } - } - } else { - for i in 0..item_len { - // SAFETY: no null buffer means every index is valid. - let s = unsafe { string_array.value_unchecked(i) }; - builder.try_append_value(&unicode_case(s, lower))?; - } - } + ColumnarValue::Array(array) => Ok(ColumnarValue::Array( + case_conversion_columnar_array(array, lower, name)?, + )), + ColumnarValue::Scalar(scalar) => Ok(ColumnarValue::Scalar( + case_conversion_scalar(scalar, lower, name)?, + )), + } +} - Ok(ColumnarValue::Array(Arc::new(builder.finish(nulls)?))) - } - other => exec_err!("Unsupported data type {other:?} for function {name}"), - }, - ColumnarValue::Scalar(scalar) => match scalar { - ScalarValue::Utf8(a) => { - let result = a.as_ref().map(|x| unicode_case(x, lower)); - Ok(ColumnarValue::Scalar(ScalarValue::Utf8(result))) - } - ScalarValue::LargeUtf8(a) => { - let result = a.as_ref().map(|x| unicode_case(x, lower)); - Ok(ColumnarValue::Scalar(ScalarValue::LargeUtf8(result))) - } - ScalarValue::Utf8View(a) => { - let result = a.as_ref().map(|x| unicode_case(x, lower)); - Ok(ColumnarValue::Scalar(ScalarValue::Utf8View(result))) +fn case_conversion_scalar( + scalar: &ScalarValue, + lower: bool, + name: &str, +) -> Result { + match scalar { + ScalarValue::Utf8(a) => { + let result = a.as_ref().map(|x| unicode_case(x, lower)); + Ok(ScalarValue::Utf8(result)) + } + ScalarValue::LargeUtf8(a) => { + let result = a.as_ref().map(|x| unicode_case(x, lower)); + Ok(ScalarValue::LargeUtf8(result)) + } + ScalarValue::Utf8View(a) => { + let result = a.as_ref().map(|x| unicode_case(x, lower)); + Ok(ScalarValue::Utf8View(result)) + } + ScalarValue::Dictionary(key_type, value) => { + let converted = case_conversion_scalar(value.as_ref(), lower, name)?; + Ok(ScalarValue::Dictionary( + key_type.clone(), + Box::new(converted), + )) + } + other => exec_err!("Unsupported data type {other:?} for function {name}"), + } +} + +fn case_conversion_columnar_array( + array: &ArrayRef, + lower: bool, + name: &str, +) -> Result { + match array.data_type() { + DataType::Utf8 => case_conversion_array::(array, lower), + DataType::LargeUtf8 => case_conversion_array::(array, lower), + DataType::Utf8View => case_conversion_utf8view(array, lower), + DataType::Dictionary(_, _) => case_conversion_dictionary(array, lower, name), + other => exec_err!("Unsupported data type {other:?} for function {name}"), + } +} + +fn case_conversion_utf8view(array: &ArrayRef, lower: bool) -> Result { + let string_array = as_string_view_array(array)?; + if string_array.is_ascii() { + return Ok(Arc::new(case_conversion_utf8view_ascii( + string_array, + lower, + ))); + } + let item_len = string_array.len(); + // Null-preserving: reuse the input null buffer as the output null buffer. + let nulls = string_array.nulls().cloned(); + let mut builder = StringViewArrayBuilder::with_capacity(item_len); + + if let Some(ref n) = nulls { + for i in 0..item_len { + if n.is_null(i) { + builder.try_append_placeholder()?; + } else { + // SAFETY: `n.is_null(i)` was false in the branch above. + let s = unsafe { string_array.value_unchecked(i) }; + builder.try_append_value(&unicode_case(s, lower))?; } - other => exec_err!("Unsupported data type {other:?} for function {name}"), - }, + } + } else { + for i in 0..item_len { + // SAFETY: no null buffer means every index is valid. + let s = unsafe { string_array.value_unchecked(i) }; + builder.try_append_value(&unicode_case(s, lower))?; + } } + + Ok(Arc::new(builder.finish(nulls)?)) +} + +fn case_conversion_dictionary( + array: &ArrayRef, + lower: bool, + name: &str, +) -> Result { + let dictionary = array.as_any_dictionary(); + let converted = case_conversion_columnar_array(dictionary.values(), lower, name)?; + Ok(dictionary.with_values(converted)) } fn case_conversion_array( diff --git a/datafusion/functions/src/string/lower.rs b/datafusion/functions/src/string/lower.rs index 57cbe1d8779f0..88f2c800e9e0c 100644 --- a/datafusion/functions/src/string/lower.rs +++ b/datafusion/functions/src/string/lower.rs @@ -21,8 +21,8 @@ use crate::string::common::to_lower; use datafusion_common::Result; use datafusion_common::types::logical_string; 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; @@ -57,9 +57,10 @@ impl LowerFunc { 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, ), } diff --git a/datafusion/functions/src/string/upper.rs b/datafusion/functions/src/string/upper.rs index c0ac90b1bc598..789ab2c046203 100644 --- a/datafusion/functions/src/string/upper.rs +++ b/datafusion/functions/src/string/upper.rs @@ -20,8 +20,8 @@ use arrow::datatypes::DataType; use datafusion_common::Result; use datafusion_common::types::logical_string; 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; @@ -56,9 +56,10 @@ impl UpperFunc { 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, ), } diff --git a/datafusion/sqllogictest/test_files/functions.slt b/datafusion/sqllogictest/test_files/functions.slt index 2b393f1a26413..98edfa189d3e3 100644 --- a/datafusion/sqllogictest/test_files/functions.slt +++ b/datafusion/sqllogictest/test_files/functions.slt @@ -468,7 +468,24 @@ Utf8View query T SELECT arrow_typeof(upper(arrow_cast(arrow_cast('foo', 'Dictionary(Int32, Utf8)'), 'Dictionary(Int32, Utf8View)'))) ---- -Utf8View +Dictionary(Int32, Utf8View) + +statement ok +CREATE TABLE upper_dictionary_test AS +SELECT arrow_cast(column1, 'Dictionary(Int32, Utf8)') AS c1 FROM (VALUES +('foo'), +(NULL), +('Bar')); + +query TT +SELECT upper(c1), arrow_typeof(upper(c1)) FROM upper_dictionary_test +---- +FOO Dictionary(Int32, Utf8) +NULL Dictionary(Int32, Utf8) +BAR Dictionary(Int32, Utf8) + +statement ok +DROP TABLE upper_dictionary_test query T SELECT btrim(' foo ') @@ -548,7 +565,24 @@ Utf8View query T SELECT arrow_typeof(lower(arrow_cast(arrow_cast('FOObar', 'Dictionary(Int32, Utf8)'), 'Dictionary(Int32, Utf8View)'))) ---- -Utf8View +Dictionary(Int32, Utf8View) + +statement ok +CREATE TABLE lower_dictionary_test AS +SELECT arrow_cast(column1, 'Dictionary(Int32, Utf8)') AS c1 FROM (VALUES +('FOO'), +(NULL), +('Bar')); + +query TT +SELECT lower(c1), arrow_typeof(lower(c1)) FROM lower_dictionary_test +---- +foo Dictionary(Int32, Utf8) +NULL Dictionary(Int32, Utf8) +bar Dictionary(Int32, Utf8) + +statement ok +DROP TABLE lower_dictionary_test query T SELECT ltrim(' foo') diff --git a/docs/source/library-user-guide/functions/adding-udfs.md b/docs/source/library-user-guide/functions/adding-udfs.md index b6021c9dbb7b4..c3a40557a006d 100644 --- a/docs/source/library-user-guide/functions/adding-udfs.md +++ b/docs/source/library-user-guide/functions/adding-udfs.md @@ -397,9 +397,9 @@ impl AsyncUpper { pub fn new() -> Self { Self { signature: Signature::new( - TypeSignature::Coercible(vec![Coercion::Exact { - desired_type: TypeSignatureClass::Native(logical_string()), - }]), + TypeSignature::Coercible(vec![Coercion::new_exact( + TypeSignatureClass::Native(logical_string()), + )]), Volatility::Volatile, ), } @@ -497,9 +497,9 @@ We can now transfer the async UDF into the normal scalar using `into_scalar_udf` # pub fn new() -> Self { # Self { # signature: Signature::new( -# TypeSignature::Coercible(vec![Coercion::Exact { -# desired_type: TypeSignatureClass::Native(logical_string()), -# }]), +# TypeSignature::Coercible(vec![Coercion::new_exact( +# TypeSignatureClass::Native(logical_string()), +# )]), # Volatility::Volatile, # ), # } diff --git a/docs/source/library-user-guide/upgrading/55.0.0.md b/docs/source/library-user-guide/upgrading/55.0.0.md index 307de722bd13e..de45ce902c5a7 100644 --- a/docs/source/library-user-guide/upgrading/55.0.0.md +++ b/docs/source/library-user-guide/upgrading/55.0.0.md @@ -82,6 +82,27 @@ will now appear as `Decimal128(NULL,10,2)`. Query result values already used human-readable decimal formatting and are unchanged. +### `Coercion` supports dictionary encoding preservation + +`datafusion_expr_common::signature::Coercion` now supports optional dictionary +encoding preservation. When enabled for `TypeSignatureClass::Native(...)` +coercions, DataFusion coerces dictionary inputs to +`Dictionary(original_key_type, coerced_value_type)` instead of materializing them +to the coerced value type. + +User-defined functions can opt in by setting dictionary encoding preservation on +the relevant coercion: + +```rust +Coercion::new_exact(TypeSignatureClass::Native(logical_string())) + .with_encoding_preservation(EncodingPreservation::dictionary()) +``` + +This changes the coerced argument type passed to the function. If a function +derives its return type from that coerced argument type, code that checks exact +result types may need to update its expectations or add an explicit cast to +materialize the result. + ### `GroupsAccumulator::merge_batch` no longer takes `opt_filter` The `opt_filter` argument has been removed from