diff --git a/datafusion/expr/src/built_in_function.rs b/datafusion/expr/src/built_in_function.rs index 196d278dc70e5..31be5f2c8a218 100644 --- 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 index e419343548065..baefb97f8f95f 100644 --- 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 index d4095a72fe3e4..b9754ece8c0bf 100644 --- 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 } - }, - ), + }), } }