diff --git a/datafusion/expr/src/type_coercion/functions.rs b/datafusion/expr/src/type_coercion/functions.rs index 65a45c078062c..37a1c10159fe8 100644 --- a/datafusion/expr/src/type_coercion/functions.rs +++ b/datafusion/expr/src/type_coercion/functions.rs @@ -33,7 +33,9 @@ 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, EncodingPreservation}; +use datafusion_expr_common::signature::{ + ArrayFunctionArgument, EncodingPreservation, TypeSignatureClass, +}; use datafusion_expr_common::type_coercion::binary::type_union_resolution; use datafusion_expr_common::{ signature::{ArrayFunctionSignature, FIXED_SIZE_LIST_WILDCARD, TIMEZONE_WILDCARD}, @@ -873,31 +875,47 @@ 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 coercion_value_type<'a>( + current_type: &'a DataType, + desired_type: &TypeSignatureClass, + ) -> &'a DataType { + if matches!(desired_type, TypeSignatureClass::Any) { + return current_type; + } + + match current_type { + DataType::Dictionary(_, value_type) => { + coercion_value_type(value_type, desired_type) + } + _ => current_type, } } fn preserve_encoding( current_type: &DataType, casted_type: DataType, + desired_type: &TypeSignatureClass, 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 + if matches!(desired_type, TypeSignatureClass::Any) { + return casted_type; + } + + match current_type { + DataType::Dictionary(key_type, value_type) => { + let casted_type = preserve_encoding( + value_type, + casted_type, + desired_type, + encoding_preservation, + ); + if encoding_preservation.preserve_dictionary() { + DataType::Dictionary(key_type.clone(), Box::new(casted_type)) + } else { + casted_type + } + } + _ => casted_type, } } @@ -905,7 +923,8 @@ fn get_valid_types( 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); + let coercion_value_type = + coercion_value_type(current_type, param.desired_type()); if param .desired_type() @@ -913,11 +932,12 @@ fn get_valid_types( { let casted_type = param .desired_type() - .default_casted_type(¤t_native_type, cast_origin)?; + .default_casted_type(¤t_native_type, coercion_value_type)?; new_types.push(preserve_encoding( current_type, casted_type, + param.desired_type(), encoding_preservation, )); } else if param @@ -928,10 +948,11 @@ 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(cast_origin)?; + default_casted_type.default_cast_for(coercion_value_type)?; new_types.push(preserve_encoding( current_type, casted_type, + param.desired_type(), encoding_preservation, )); } else { @@ -1851,16 +1872,33 @@ mod tests { ))?; assert_eq!(vec![DataType::Int64], output); - // Dictionary gets passed through if we use TypeSignatureClass apart from Native - let output = dictionary_input(Coercion::new_exact(TypeSignatureClass::Integer))?; + // Any always preserves the original physical type + let output = dictionary_input(Coercion::new_exact(TypeSignatureClass::Any))?; assert_eq!(vec![dictionary.clone()], output); + let output = dictionary_input( + Coercion::new_exact(TypeSignatureClass::Any) + .with_encoding_preservation(EncodingPreservation::dictionary()), + )?; + assert_eq!(vec![dictionary.clone()], output); + + // Typed non-Native classes materialize dictionaries by default + let output = dictionary_input(Coercion::new_exact(TypeSignatureClass::Integer))?; + assert_eq!(vec![DataType::Int64], output); + let output = dictionary_input(Coercion::new_implicit( TypeSignatureClass::Integer, vec![], NativeType::Int64, ))?; - assert_eq!(vec![dictionary.clone()], output); + assert_eq!(vec![DataType::Int64], output); + + // Typed non-Native classes preserve dictionaries only when requested + let output = dictionary_input( + Coercion::new_exact(TypeSignatureClass::Integer) + .with_encoding_preservation(EncodingPreservation::dictionary()), + )?; + assert_eq!(vec![dictionary], output); Ok(()) } @@ -1936,7 +1974,7 @@ mod tests { Box::new(DataType::Int64), )] ); - // Contrast: without encoding_preservation, non-Native already passes through + // Without encoding_preservation, non-Native classes materialize dictionaries assert_eq!( dictionary_input( DataType::Int32, @@ -1946,12 +1984,9 @@ mod tests { NativeType::Int64, ), )?, - vec![DataType::Dictionary( - Box::new(DataType::Int8), - Box::new(DataType::Int32), - )] + vec![DataType::Int32] ); - // With encoding_preservation, same result — no difference for non-Native + // With encoding_preservation, non-Native classes preserve dictionaries assert_eq!( dictionary_input( DataType::Int32, @@ -1971,6 +2006,66 @@ mod tests { Ok(()) } + #[test] + fn test_coercible_nested_dictionary() -> Result<()> { + let nested_dictionary = DataType::Dictionary( + Box::new(DataType::Int8), + Box::new(DataType::Dictionary( + Box::new(DataType::Int16), + Box::new(DataType::Int32), + )), + ); + let nested_dictionary_input = |coercion| -> Result> { + fields_with_udf( + &[Field::new("field", nested_dictionary.clone(), true).into()], + &MockUdf(Signature::coercible(vec![coercion], Volatility::Immutable)), + ) + .map(|v| v.into_iter().map(|f| f.data_type().clone()).collect()) + }; + + // Without preservation, recursively unwrap dictionaries to the unchanged leaf. + let output = + nested_dictionary_input(Coercion::new_exact(TypeSignatureClass::Integer))?; + assert_eq!(vec![DataType::Int32], output); + + // With preservation, restore the complete dictionary stack around the leaf. + let output = nested_dictionary_input( + Coercion::new_exact(TypeSignatureClass::Integer) + .with_encoding_preservation(EncodingPreservation::dictionary()), + )?; + assert_eq!(vec![nested_dictionary.clone()], output); + + let int64_coercion = || { + Coercion::new_implicit( + TypeSignatureClass::Native(logical_int64()), + vec![TypeSignatureClass::Integer], + NativeType::Int64, + ) + }; + + // Without preservation, materialize the coerced leaf type. + let output = nested_dictionary_input(int64_coercion())?; + assert_eq!(vec![DataType::Int64], output); + + // With preservation, restore the complete dictionary stack around the coerced leaf. + let output = nested_dictionary_input( + int64_coercion() + .with_encoding_preservation(EncodingPreservation::dictionary()), + )?; + assert_eq!( + vec![DataType::Dictionary( + Box::new(DataType::Int8), + Box::new(DataType::Dictionary( + Box::new(DataType::Int16), + Box::new(DataType::Int64), + )), + )], + output + ); + + Ok(()) + } + #[test] fn test_coercible_run_end_encoded() -> Result<()> { let run_end_encoded = DataType::RunEndEncoded( diff --git a/datafusion/spark/src/function/bitmap/bitmap_count.rs b/datafusion/spark/src/function/bitmap/bitmap_count.rs index 89bea101afbe7..18d584868830b 100644 --- a/datafusion/spark/src/function/bitmap/bitmap_count.rs +++ b/datafusion/spark/src/function/bitmap/bitmap_count.rs @@ -28,8 +28,8 @@ use arrow::datatypes::{DataType, FieldRef, Int8Type, Int16Type, Int32Type, Int64 use datafusion_common::utils::take_function_args; use datafusion_common::{Result, internal_err}; use datafusion_expr::{ - Coercion, ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, - TypeSignatureClass, Volatility, + Coercion, ColumnarValue, EncodingPreservation, ScalarFunctionArgs, ScalarUDFImpl, + Signature, TypeSignatureClass, Volatility, }; use datafusion_functions::downcast_arg; use datafusion_functions::utils::make_scalar_function; @@ -49,7 +49,10 @@ impl BitmapCount { pub fn new() -> Self { Self { signature: Signature::coercible( - vec![Coercion::new_exact(TypeSignatureClass::Binary)], + vec![ + Coercion::new_exact(TypeSignatureClass::Binary) + .with_encoding_preservation(EncodingPreservation::dictionary()), + ], Volatility::Immutable, ), } diff --git a/datafusion/spark/src/function/math/hex.rs b/datafusion/spark/src/function/math/hex.rs index 55c9cda63c888..c7b82d53735a3 100644 --- a/datafusion/spark/src/function/math/hex.rs +++ b/datafusion/spark/src/function/math/hex.rs @@ -34,8 +34,8 @@ use datafusion_common::{ exec_datafusion_err, exec_err, }; use datafusion_expr::{ - Coercion, ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, - TypeSignatureClass, Volatility, + Coercion, ColumnarValue, EncodingPreservation, ScalarFunctionArgs, ScalarUDFImpl, + Signature, TypeSignature, TypeSignatureClass, Volatility, }; /// #[derive(Debug, PartialEq, Eq, Hash)] @@ -60,7 +60,8 @@ impl SparkHex { let string = Coercion::new_exact(TypeSignatureClass::Native(logical_string())); - let binary = Coercion::new_exact(TypeSignatureClass::Binary); + let binary = Coercion::new_exact(TypeSignatureClass::Binary) + .with_encoding_preservation(EncodingPreservation::dictionary()); let variants = vec![ // accepts numeric types diff --git a/datafusion/sqllogictest/test_files/expr.slt b/datafusion/sqllogictest/test_files/expr.slt index 7e15b48a0d824..c7c997f330d93 100644 --- a/datafusion/sqllogictest/test_files/expr.slt +++ b/datafusion/sqllogictest/test_files/expr.slt @@ -850,6 +850,15 @@ SELECT to_hex(0) ---- 0 +query T +SELECT to_hex(arrow_cast(a, 'Dictionary(Int32, Int64)')) +FROM (VALUES (0), (10), (255), (NULL)) AS t(a) +---- +0 +a +ff +NULL + # negative values (two's complement encoding) query T SELECT to_hex(-1) diff --git a/datafusion/sqllogictest/test_files/spark/bitmap/bitmap_count.slt b/datafusion/sqllogictest/test_files/spark/bitmap/bitmap_count.slt index 39dca512226b2..3ac5337cd7fd5 100644 --- a/datafusion/sqllogictest/test_files/spark/bitmap/bitmap_count.slt +++ b/datafusion/sqllogictest/test_files/spark/bitmap/bitmap_count.slt @@ -68,6 +68,21 @@ SELECT bitmap_count(arrow_cast(a, 'Dictionary(Int32, Binary)')) FROM (VALUES (X' 16 NULL +# The CAST to Dictionary below comes from the explicit arrow_cast. There must +# not be an additional outer CAST(... AS Binary) before bitmap_count. +query TT +EXPLAIN SELECT bitmap_count(arrow_cast(a, 'Dictionary(Int32, Binary)')) +FROM (VALUES (X'1010'), (X'0AB0'), (X'FFFF'), (NULL)) AS t(a); +---- +logical_plan +01)Projection: bitmap_count(CAST(t.a AS Dictionary(Int32, Binary))) AS bitmap_count(arrow_cast(t.a,Utf8("Dictionary(Int32, Binary)"))) +02)--SubqueryAlias: t +03)----Projection: column1 AS a +04)------Values: (Binary("16,16")), (Binary("10,176")), (Binary("255,255")), (Binary(NULL)) +physical_plan +01)ProjectionExec: expr=[bitmap_count(CAST(column1@0 AS Dictionary(Int32, Binary))) as bitmap_count(arrow_cast(t.a,Utf8("Dictionary(Int32, Binary)")))] +02)--DataSourceExec: partitions=1, partition_sizes=[1] + query I SELECT bitmap_count(arrow_cast(a, 'Dictionary(Int8, Binary)')) FROM (VALUES (X'1010'), (X'0AB0'), (X'FFFF'), (NULL)) AS t(a); ---- 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 181e0e0b7f266..3dac6422056d8 100644 --- a/docs/source/library-user-guide/upgrading/55.0.0.md +++ b/docs/source/library-user-guide/upgrading/55.0.0.md @@ -168,8 +168,10 @@ 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 +encoding preservation. Typed coercions materialize dictionary inputs by +default, including both `TypeSignatureClass::Native(...)` and broader classes +such as `Integer`, `Numeric`, and `Binary`. When preservation is enabled, +DataFusion instead coerces dictionary inputs to `Dictionary(original_key_type, coerced_value_type)` instead of materializing them to the coerced value type. @@ -186,6 +188,11 @@ 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. +This changes the previous behavior of typed non-native classes such as +`Integer` and `Binary`, which retained the physical dictionary type by default. +UDFs relying on that behavior must now explicitly enable dictionary +preservation. `TypeSignatureClass::Any` is unaffected. + ### `GroupsAccumulator::merge_batch` no longer takes `opt_filter` The `opt_filter` argument has been removed from