Coerce Dictionary types for scalar functions

viirya · Apr 14, 2024 · 531bd1a · 531bd1a
1 parent 7f497b3
commit 531bd1a
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 6 deletions.
diff --git a/datafusion/expr/src/built_in_function.rs b/datafusion/expr/src/built_in_function.rs
@@ -635,4 +635,19 @@ mod tests {
             .unwrap();
         assert_eq!(return_type, DataType::Date32);
     }
+
+    #[test]
+    fn test_coalesce_return_types_dictionary() {
+        let coalesce = BuiltinScalarFunction::Coalesce;
+        let return_type = coalesce
+            .return_type(&[
+                DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
+                DataType::Utf8,
+            ])
+            .unwrap();
+        assert_eq!(
+            return_type,
+            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8))
+        );
+    }
 }
diff --git a/datafusion/expr/src/type_coercion/binary.rs b/datafusion/expr/src/type_coercion/binary.rs
@@ -632,7 +632,7 @@ fn both_numeric_or_null_and_numeric(lhs_type: &DataType, rhs_type: &DataType) ->
 ///
 /// Not all operators support dictionaries, if `preserve_dictionaries` is true
 /// dictionaries will be preserved if possible
-fn dictionary_coercion(
+pub(crate) fn dictionary_coercion(
     lhs_type: &DataType,
     rhs_type: &DataType,
     preserve_dictionaries: bool,

diff --git a/datafusion/expr/src/type_coercion/functions.rs b/datafusion/expr/src/type_coercion/functions.rs
@@ -28,7 +28,9 @@ use arrow::{
 use datafusion_common::utils::{coerced_fixed_size_list_to_list, list_ndims};
 use datafusion_common::{internal_datafusion_err, internal_err, plan_err, Result};
 
-use super::binary::{comparison_binary_numeric_coercion, comparison_coercion};
+use super::binary::{
+    comparison_binary_numeric_coercion, comparison_coercion, dictionary_coercion,
+};
 
 /// Performs type coercion for function arguments.
 ///
@@ -435,15 +437,15 @@ fn coerced_from<'a>(
         // Note that not all rules in `comparison_coercion` can be reused here.
         // For example, all numeric types can be coerced into Utf8 for comparison,
         // but not for function arguments.
-        _ => comparison_binary_numeric_coercion(type_into, type_from).and_then(
-            |coerced_type| {
+        _ => comparison_binary_numeric_coercion(type_into, type_from)
+            .or_else(|| dictionary_coercion(type_into, type_from, true))
+            .and_then(|coerced_type| {
                 if *type_into == coerced_type {
                     Some(coerced_type)
                 } else {
                     None
                 }
-            },
-        ),
+            }),
     }
 }