[branch-56] fix: infer UNION output types before analysis (backport #27139)

datafusion/core/tests/sql/mod.rs

@@ -61,6 +61,7 @@ mod sql_api;
 mod tpcds_planning;
 mod union_comparison;
 mod union_nullable;
+mod union_schema;
 mod unparser;
 
 async fn register_aggregate_csv_by_sql(ctx: &SessionContext) {

datafusion/core/tests/sql/union_nullable.rs

@@ -61,12 +61,8 @@ async fn union_all_mismatched_nullable(
     .await
 }
 
-/// The schema DataFusion actually commits to for a query: the logical plan
-/// after the `Analyzer` (which includes the `UNION` nullability/type
-/// coercion this test targets) and `Optimizer` have run. `DataFrame::schema`
-/// alone is not enough here -- it reflects the raw, pre-`Analyzer` plan (see
-/// `SessionState::create_logical_plan`), which for a `UNION` still has the
-/// first leg's un-coerced type.
+/// This test checks the schema after the Analyzer and Optimizer run,
+/// because they can change nullability.
 fn analyzed_schema(df: &DataFrame) -> Result<Schema> {
     Ok(df
         .clone()

datafusion/core/tests/sql/union_schema.rs

@@ -0,0 +1,138 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
+// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
+
+//! The pre-analysis UNION schema must describe the output types (#27023).
+
+use arrow::datatypes::{DataType, TimeUnit};
+use datafusion::prelude::*;
+use datafusion_common::{Result, ScalarValue};
+
+fn assert_union_schema(df: &DataFrame, expected: &[(&str, DataType)]) {
+    let fields = df.schema().fields();
+    assert_eq!(fields.len(), expected.len());
+    for (field, (name, data_type)) in fields.iter().zip(expected) {
+        assert_eq!(field.name(), name);
+        assert_eq!(field.data_type(), data_type);
+    }
+}
+
+#[tokio::test]
+async fn union_infers_types_before_analysis() -> Result<()> {
+    let ctx = SessionContext::new();
+    let cases = [
+        (
+            r#"arrow_cast(arrow_cast(1, 'Timestamp(Microsecond, None)'), 'Timestamp(Microsecond, Some("UTC"))')"#,
+            "CAST(NULL AS TIMESTAMP)",
+            DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into())),
+        ),
+        (
+            "arrow_cast('a', 'Utf8')",
+            "arrow_cast('b', 'LargeUtf8')",
+            DataType::LargeUtf8,
+        ),
+        (
+            "arrow_cast('a', 'Utf8')",
+            "arrow_cast('b', 'Utf8View')",
+            DataType::Utf8View,
+        ),
+        ("CAST(1 AS SMALLINT)", "2.5", DataType::Float64),
+        (
+            "CASE WHEN false THEN 1 ELSE 2.5 END",
+            "NULL",
+            DataType::Float64,
+        ),
+        (
+            "CASE WHEN false THEN 1 WHEN true THEN 2.5 END",
+            "NULL",
+            DataType::Float64,
+        ),
+    ];
+    for (left, right, expected) in cases {
+        for (left, right) in [(left, right), (right, left)] {
+            for operation in ["ALL", "DISTINCT", "ALL BY NAME", "BY NAME"] {
+                let sql = format!(
+                    "SELECT x FROM (SELECT {left} AS x UNION {operation} SELECT {right} AS x) u"
+                );
+                assert_union_schema(&ctx.sql(&sql).await?, &[("x", expected.clone())]);
+            }
+        }
+    }
+    Ok(())
+}
+
+#[tokio::test]
+async fn union_infers_dictionary_types_before_analysis() -> Result<()> {
+    let ctx = SessionContext::new();
+    for (sql, expected) in [
+        (
+            "SELECT NULL AS x UNION ALL SELECT arrow_cast('a', 'Dictionary(Int32, Utf8)') AS x",
+            DataType::Utf8,
+        ),
+        (
+            "SELECT CAST(1 AS SMALLINT) AS x UNION ALL SELECT arrow_cast(2, 'Dictionary(Int32, Int32)') AS x",
+            DataType::Int32,
+        ),
+    ] {
+        assert_union_schema(&ctx.sql(sql).await?, &[("x", expected)]);
+    }
+    for sql in [
+        "SELECT 1 AS y UNION ALL BY NAME SELECT arrow_cast('a', 'Dictionary(Int32, Utf8)') AS x",
+        "SELECT arrow_cast('a', 'Dictionary(Int32, Utf8)') AS x UNION ALL BY NAME SELECT 1 AS y",
+    ] {
+        let df = ctx.sql(sql).await?.select_columns(&["x", "y"])?;
+        assert_union_schema(&df, &[("x", DataType::Utf8), ("y", DataType::Int64)]);
+    }
+    Ok(())
+}
+
+#[tokio::test]
+async fn union_dataframe_and_by_name_schemas() -> Result<()> {
+    let ctx = SessionContext::new();
+    for variant in 0..4 {
+        let left = ctx.sql("SELECT CAST(1 AS SMALLINT) AS x").await?;
+        let right = ctx.sql("SELECT 2.5 AS x").await?;
+        let df = match variant {
+            0 => left.union(right)?,
+            1 => left.union_distinct(right)?,
+            2 => left.union_by_name(right)?,
+            _ => left.union_by_name_distinct(right)?,
+        };
+        assert_union_schema(&df.select_columns(&["x"])?, &[("x", DataType::Float64)]);
+    }
+    for sql in [
+        "SELECT CAST(1 AS SMALLINT) AS x, 1 AS y UNION ALL BY NAME SELECT 2 AS y, 2.5 AS x",
+        "SELECT CAST(1 AS SMALLINT) AS x UNION ALL BY NAME SELECT 2.5 AS x, 2 AS y",
+    ] {
+        assert_union_schema(
+            &ctx.sql(sql).await?,
+            &[("x", DataType::Float64), ("y", DataType::Int64)],
+        );
+    }
+    let df = ctx
+        .sql(
+            "SELECT CAST(NULL AS SMALLINT) AS x WHERE false \
+         UNION ALL SELECT CAST(NULL AS BIGINT) AS x WHERE false",
+        )
+        .await?;
+    assert_union_schema(&df, &[("x", DataType::Int64)]);
+    let df = ctx
+        .sql(
+            "SELECT max(x) AS x FROM (SELECT 1 AS x UNION ALL SELECT 2.5 AS x \
+         UNION ALL SELECT 3 AS x) u",
+        )
+        .await?;
+    assert_union_schema(&df, &[("x", DataType::Float64)]);
+    Ok(())
+}
+
+#[tokio::test]
+async fn union_bound_parameters_keep_inferred_type() -> Result<()> {
+    let ctx = SessionContext::new();
+    let df = ctx
+        .sql("SELECT CAST($1 AS INT) AS x UNION ALL SELECT 2.5 AS x")
+        .await?
+        .with_param_values(vec![ScalarValue::Int32(Some(1))])?;
+    assert_union_schema(&df, &[("x", DataType::Float64)]);
+    Ok(())
+}

datafusion/expr/src/expr_schema.rs

@@ -13,6 +13,7 @@ use crate::expr::{FieldMetadata, LambdaVariable};
 use crate::higher_order_function::HigherOrderReturnFieldArgs;
 use crate::type_coercion::functions::value_fields_with_higher_order_udf_and_lambdas;
 use crate::type_coercion::functions::{UDFCoercionExt, fields_with_udf};
+use crate::type_coercion::other::get_coerce_type_for_case_expression;
 use crate::udf::ReturnFieldArgs;
 use crate::{
     LogicalPlan, Operator, Projection, Subquery, WindowFunctionDefinition, utils,
@@ -33,6 +34,10 @@ use std::sync::Arc;
 /// Trait to allow expr to typable with respect to a schema
 pub trait ExprSchemable {
     /// Given a schema, return the type of the expr
+    ///
+    /// `CASE` expressions infer a common type from all `THEN` and `ELSE`
+    /// branches without inserting casts. Incompatible branches retain a
+    /// provisional type until analysis.
     fn get_type(&self, schema: &dyn ExprSchema) -> Result<DataType>;
 
     /// Given a schema, return the nullability of the expr
@@ -177,15 +182,28 @@ impl ExprSchemable for Expr {
             Expr::ScalarVariable(field, _) => Ok(field.data_type().clone()),
             Expr::Literal(l, _) => Ok(l.data_type()),
             Expr::Case(case) => {
-                for (_, then_expr) in &case.when_then_expr {
-                    let then_type = then_expr.get_type(schema)?;
-                    if !then_type.is_null() {
-                        return Ok(then_type);
-                    }
-                }
-                case.else_expr
+                let then_types = case
+                    .when_then_expr
+                    .iter()
+                    .map(|(_, expr)| expr.get_type(schema))
+                    .collect::<Result<Vec<_>>>()?;
+                let else_type = case
+                    .else_expr
                     .as_ref()
-                    .map_or(Ok(DataType::Null), |e| e.get_type(schema))
+                    .map(|expr| expr.get_type(schema))
+                    .transpose()?;
+                // If the branches have no common type, keep the old provisional
+                // type. The analyzer reports the error.
+                let common_type =
+                    get_coerce_type_for_case_expression(&then_types, else_type.as_ref());
+                let result_type = common_type.unwrap_or_else(|| {
+                    then_types
+                        .into_iter()
+                        .find(|t| !t.is_null())
+                        .or(else_type)
+                        .unwrap_or(DataType::Null)
+                });
+                Ok(result_type)
             }
             Expr::Cast(Cast { field, .. }) | Expr::TryCast(TryCast { field, .. }) => {
                 Ok(field.data_type().clone())
@@ -1684,4 +1702,33 @@ mod tests {
             }
         }
     }
+
+    #[test]
+    fn case_infers_all_result_types_without_casting() -> Result<()> {
+        use crate::expr::Case;
+        let schema = DFSchema::empty();
+        let nested = Expr::Case(Case::new(
+            None,
+            vec![(Box::new(lit(false)), Box::new(lit(1_i64)))],
+            Some(Box::new(lit(2.5))),
+        ));
+        for (then_values, else_value, expected) in [
+            (vec![lit(1_i64)], Some(lit(2.5)), DataType::Float64),
+            (vec![lit(1_i64), lit(2.5)], None, DataType::Float64),
+            (vec![lit(2.5), lit(1_i64)], None, DataType::Float64),
+            (vec![lit(2.5)], Some(lit(1_i64)), DataType::Float64),
+            (vec![lit(3_i64)], Some(nested), DataType::Float64),
+            (vec![lit(ScalarValue::Null)], None, DataType::Null),
+            // Incompatible raw types retain a provisional type for the analyzer.
+            (vec![lit(true)], Some(lit("a")), DataType::Boolean),
+        ] {
+            let when_then = then_values
+                .into_iter()
+                .map(|value| (Box::new(lit(true)), Box::new(value)))
+                .collect();
+            let expr = Expr::Case(Case::new(None, when_then, else_value.map(Box::new)));
+            assert_eq!(expr.get_type(&schema)?, expected);
+        }
+        Ok(())
+    }
 }

datafusion/expr/src/logical_plan/builder.rs

@@ -2189,8 +2189,8 @@ pub fn validate_unique_names<'a>(
 
 /// Union two [`LogicalPlan`]s.
 ///
-/// Constructs the UNION plan, but does not perform type-coercion. Therefore the
-/// subtree expressions will not be properly typed until the optimizer pass.
+/// Infers common output types, but does not insert casts into the inputs.
+/// The analyzer must still resolve input expressions and insert any required casts.
 ///
 /// If a properly typed UNION plan is needed, refer to [`TypeCoercionRewriter::coerce_union`]
 /// or alternatively, merge the union input schema using [`coerce_union_schema`] and