fix: align dictionary coercion across typed signatures

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,51 +875,69 @@ 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,
                 }
             }
 
             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);
+                let coercion_value_type =
+                    coercion_value_type(current_type, param.desired_type());
 
                 if param
                     .desired_type()
                     .matches_native_type(&current_native_type)
                 {
                     let casted_type = param
                         .desired_type()
-                        .default_casted_type(&current_native_type, cast_origin)?;
+                        .default_casted_type(&current_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<Vec<DataType>> {
+            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(

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,
             ),
         }

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,
 };
 /// <https://spark.apache.org/docs/latest/api/sql/index.html#hex>
 #[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

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)

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);
 ----

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