fix: Coerce CASE branch types to a common type

datafusion/core/src/logical_plan/expr_schema.rs

@@ -20,7 +20,7 @@ use crate::logical_expr::{aggregate_function, function, window_function};
 use arrow::compute::can_cast_types;
 use arrow::datatypes::DataType;
 use datafusion_common::{DFField, DFSchema, DataFusionError, ExprSchema, Result};
-use datafusion_expr::binary_rule::binary_operator_data_type;
+use datafusion_expr::binary_rule::{binary_operator_data_type, case_expression_coercion};
 use datafusion_physical_expr::field_util::get_indexed_field;
 
 /// trait to allow expr to typable with respect to a schema
@@ -60,7 +60,25 @@ impl ExprSchemable for Expr {
             Expr::OuterColumn(ty, _) => Ok(ty.clone()),
             Expr::ScalarVariable(ty, _) => Ok(ty.clone()),
             Expr::Literal(l) => Ok(l.get_datatype()),
-            Expr::Case { when_then_expr, .. } => when_then_expr[0].1.get_type(schema),
+            Expr::Case {
+                when_then_expr,
+                else_expr,
+                ..
+            } => {
+                let mut branch_types = when_then_expr
+                    .iter()
+                    .map(|(_, then_expr)| then_expr.get_type(schema))
+                    .collect::<Result<Vec<_>>>()?;
+                if let Some(else_expr) = else_expr {
+                    branch_types.push(else_expr.get_type(schema)?);
+                }
+                case_expression_coercion(&branch_types).ok_or_else(|| {
+                    DataFusionError::Plan(format!(
+                        "CASE branches have no common type to coerce the results to: {:?}",
+                        branch_types
+                    ))
+                })
+            }
             Expr::Cast { data_type, .. } | Expr::TryCast { data_type, .. } => {
                 Ok(data_type.clone())
             }

datafusion/expr/src/binary_rule.rs

@@ -196,6 +196,44 @@ pub fn comparison_eq_coercion(
         .or_else(|| string_boolean_coercion(lhs_type, rhs_type))
 }
 
+/// Coerce the THEN/ELSE branch types of a `CASE` expression to a single common type.
+///
+/// `Null` branches are ignored, so `CASE WHEN c THEN NULL ELSE 1 END` resolves to
+/// `Int64`; when every branch is `Null` the result is `Null`. Returns `None` when the
+/// branches have no common type.
+pub fn case_expression_coercion<'a>(
+    branch_types: impl IntoIterator<Item = &'a DataType>,
+) -> Option<DataType> {
+    let mut result = DataType::Null;
+    for branch_type in branch_types {
+        if matches!(branch_type, DataType::Null) {
+            continue;
+        }
+        result = match result {
+            DataType::Null => branch_type.clone(),
+            current => case_branch_coercion(&current, branch_type)?,
+        };
+    }
+    Some(result)
+}
+
+/// Coercion rules for a pair of `CASE` branch types.
+///
+/// This is the equality coercion without the string/boolean and string/numeric rules that
+/// narrow a string down to the other side: a `CASE` result mixing a string with a boolean
+/// or a number widens to the string instead.
+fn case_branch_coercion(lhs_type: &DataType, rhs_type: &DataType) -> Option<DataType> {
+    if lhs_type == rhs_type && !is_dictionary(lhs_type) {
+        return Some(lhs_type.clone());
+    }
+    comparison_binary_numeric_coercion(lhs_type, rhs_type)
+        .or_else(|| dictionary_coercion(lhs_type, rhs_type))
+        .or_else(|| temporal_coercion(lhs_type, rhs_type))
+        .or_else(|| string_coercion(lhs_type, rhs_type))
+        .or_else(|| string_numeric_coercion(lhs_type, rhs_type))
+        .or_else(|| string_boolean_coercion(lhs_type, rhs_type))
+}
+
 // NOTE: NULL hack!
 fn string_numeric_coercion(lhs_type: &DataType, rhs_type: &DataType) -> Option<DataType> {
     use arrow::datatypes::DataType::*;