(arrow-datafusion) branch main updated: Coerce Dictionary types for scalar functions (#10077)

[alamb]

This is an automated email from the ASF dual-hosted git repository.

alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git


The following commit(s) were added to refs/heads/main by this push:
     new 6ca9d104b8 Coerce Dictionary types for scalar functions (#10077)
6ca9d104b8 is described below

commit 6ca9d104b8092ab7abac5c3603601cdb03fa07fd
Author: Liang-Chi Hsieh <vii...@gmail.com>
AuthorDate: Mon Apr 15 12:03:44 2024 -0700

    Coerce Dictionary types for scalar functions (#10077)
    
    * Coerce Dictionary types for scalar functions
    
    * Fix
    
    * Fix format
    
    * Add test
---
 datafusion/expr/src/built_in_function.rs       | 15 +++++++++++++++
 datafusion/expr/src/type_coercion/functions.rs | 25 +++++++++++++++++++++++--
 2 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/datafusion/expr/src/built_in_function.rs 
b/datafusion/expr/src/built_in_function.rs
index 5bfec00ea3..7ec544a57e 100644
--- a/datafusion/expr/src/built_in_function.rs
+++ b/datafusion/expr/src/built_in_function.rs
@@ -277,4 +277,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/functions.rs 
b/datafusion/expr/src/type_coercion/functions.rs
index 34b607d088..37eeb7d464 100644
--- a/datafusion/expr/src/type_coercion/functions.rs
+++ b/datafusion/expr/src/type_coercion/functions.rs
@@ -314,8 +314,13 @@ fn coerced_from<'a>(
     // match Dictionary first
     match (type_into, type_from) {
         // coerced dictionary first
-        (cur_type, Dictionary(_, value_type)) | (Dictionary(_, value_type), 
cur_type)
-            if coerced_from(cur_type, value_type).is_some() =>
+        (_, Dictionary(_, value_type))
+            if coerced_from(type_into, value_type).is_some() =>
+        {
+            Some(type_into.clone())
+        }
+        (Dictionary(_, value_type), _)
+            if coerced_from(value_type, type_from).is_some() =>
         {
             Some(type_into.clone())
         }
@@ -624,4 +629,20 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn test_coerced_from_dictionary() {
+        let type_into =
+            DataType::Dictionary(Box::new(DataType::Int32), 
Box::new(DataType::UInt32));
+        let type_from = DataType::Int64;
+        assert_eq!(coerced_from(&type_into, &type_from), None);
+
+        let type_from =
+            DataType::Dictionary(Box::new(DataType::Int32), 
Box::new(DataType::UInt32));
+        let type_into = DataType::Int64;
+        assert_eq!(
+            coerced_from(&type_into, &type_from),
+            Some(type_into.clone())
+        );
+    }
 }