Coerce types for all union children plans when eliminating nesting (a…

…pache#11386)
xinlifoobar · Jul 17, 2024 · 1b40f84 · 1b40f84
1 parent 692a21d
commit 1b40f84
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 6 deletions.
diff --git a/datafusion/optimizer/src/eliminate_nested_union.rs b/datafusion/optimizer/src/eliminate_nested_union.rs
@@ -60,7 +60,8 @@ impl OptimizerRule for EliminateNestedUnion {
                 let inputs = inputs
                     .into_iter()
                     .flat_map(extract_plans_from_union)
-                    .collect::<Vec<_>>();
+                    .map(|plan| coerce_plan_expr_for_schema(&plan, &schema))
+                    .collect::<Result<Vec<_>>>()?;
 
                 Ok(Transformed::yes(LogicalPlan::Union(Union {
                     inputs: inputs.into_iter().map(Arc::new).collect_vec(),
@@ -74,7 +75,8 @@ impl OptimizerRule for EliminateNestedUnion {
                             .into_iter()
                             .map(extract_plan_from_distinct)
                             .flat_map(extract_plans_from_union)
-                            .collect::<Vec<_>>();
+                            .map(|plan| coerce_plan_expr_for_schema(&plan, &schema))
+                            .collect::<Result<Vec<_>>>()?;
 
                         Ok(Transformed::yes(LogicalPlan::Distinct(Distinct::All(
                             Arc::new(LogicalPlan::Union(Union {
@@ -95,10 +97,9 @@ impl OptimizerRule for EliminateNestedUnion {
 
 fn extract_plans_from_union(plan: Arc<LogicalPlan>) -> Vec<LogicalPlan> {
     match unwrap_arc(plan) {
-        LogicalPlan::Union(Union { inputs, schema }) => inputs
-            .into_iter()
-            .map(|plan| coerce_plan_expr_for_schema(&plan, &schema).unwrap())
-            .collect::<Vec<_>>(),
+        LogicalPlan::Union(Union { inputs, .. }) => {
+            inputs.into_iter().map(unwrap_arc).collect::<Vec<_>>()
+        }
         plan => vec![plan],
     }
 }

diff --git a/datafusion/sqllogictest/test_files/union.slt b/datafusion/sqllogictest/test_files/union.slt
@@ -135,6 +135,21 @@ SELECT SUM(d) FROM (
 ----
 5
 
+# three way union with aggregate and type coercion
+query II rowsort
+SELECT c1, SUM(c2) FROM (
+    SELECT 1 as c1, 1::int as c2
+    UNION
+    SELECT 2 as c1, 2::int as c2
+    UNION
+    SELECT 3 as c1, COALESCE(3::int, 0) as c2
+) as a
+GROUP BY c1
+----
+1 1
+2 2
+3 3
+
 # union_all_with_count
 statement ok
 CREATE table t as SELECT 1 as a