diff --git a/datafusion/core/src/logical_plan/expr_rewriter.rs b/datafusion/core/src/logical_plan/expr_rewriter.rs index a35eb85a0f19e..09ca8d6ad8605 100644 --- a/datafusion/core/src/logical_plan/expr_rewriter.rs +++ b/datafusion/core/src/logical_plan/expr_rewriter.rs @@ -619,6 +619,13 @@ pub fn coerce_plan_expr_for_schema( Box::new(e.clone().cast_to(new_type, input.schema())?), alias.clone(), )), + // Projection expressions must be resolved against the input schema: + // the projection's own schema may have had its qualifiers replaced + // (e.g. by `project_with_column_index_alias` during UNION planning), + // so qualified columns in `expr` would no longer resolve against it + (LogicalPlan::Projection(Projection { input, .. }), _) => { + expr.cast_to(new_type, input.schema()) + } _ => expr.cast_to(new_type, plan.schema()), } } else { diff --git a/datafusion/core/tests/sql/union.rs b/datafusion/core/tests/sql/union.rs index 54807b3b2da66..5a42011e3dd2c 100644 --- a/datafusion/core/tests/sql/union.rs +++ b/datafusion/core/tests/sql/union.rs @@ -75,3 +75,34 @@ async fn union_all_with_aggregate() -> Result<()> { assert_batches_eq!(expected, &actual); Ok(()) } + +#[tokio::test] +async fn union_all_ctes_with_type_coercion() -> Result<()> { + let ctx = SessionContext::new(); + // Type mismatch must enter at the second UNION level: the first UNION + // replaces its input projections' schemas with the unqualified union + // schema, and the second UNION's coercion used to fail resolving the + // still-qualified columns against them + let sql = "WITH \ + a AS (SELECT CAST(1 AS bigint) AS t), \ + b AS (SELECT CAST(2 AS bigint) AS t), \ + c AS (SELECT 3.5 AS t) \ + SELECT 'A' AS l, t FROM a \ + UNION ALL \ + SELECT 'B' AS l, t FROM b \ + UNION ALL \ + SELECT 'C' AS l, t FROM c"; + let actual = execute_to_batches(&ctx, sql).await; + #[rustfmt::skip] + let expected = [ + "+---+-----+", + "| l | t |", + "+---+-----+", + "| A | 1 |", + "| B | 2 |", + "| C | 3.5 |", + "+---+-----+", + ]; + assert_batches_eq!(expected, &actual); + Ok(()) +}