Fix UNION ALL aliasing (#6373)

* Fix UNION ALL aliasing * fmt * fmt
apache · May 18, 2023 · 0536219 · 0536219
1 parent 7c6b41a
commit 0536219
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 3 deletions.
diff --git a/datafusion/core/tests/sqllogictests/test_files/union.slt b/datafusion/core/tests/sqllogictests/test_files/union.slt
@@ -446,3 +446,30 @@ drop table t1
 
 statement ok
 drop table t2
+
+# test UNION ALL aliases correctly with all aliased
+query TT
+explain select 1 a group by a union all select 2 b union all select 3 c
+----
+logical_plan
+Union
+--Projection: Int64(1) AS a
+----Aggregate: groupBy=[[Int64(1)]], aggr=[[]]
+------EmptyRelation
+--Projection: Int64(2) AS a
+----EmptyRelation
+--Projection: Int64(3) AS a
+----EmptyRelation
+physical_plan
+UnionExec
+--ProjectionExec: expr=[Int64(1)@0 as a]
+----AggregateExec: mode=FinalPartitioned, gby=[Int64(1)@0 as Int64(1)], aggr=[]
+------CoalesceBatchesExec: target_batch_size=8192
+--------RepartitionExec: partitioning=Hash([Column { name: "Int64(1)", index: 0 }], 4), input_partitions=4
+----------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
+------------AggregateExec: mode=Partial, gby=[1 as Int64(1)], aggr=[]
+--------------EmptyExec: produce_one_row=true
+--ProjectionExec: expr=[2 as a]
+----EmptyExec: produce_one_row=true
+--ProjectionExec: expr=[3 as a]
+----EmptyExec: produce_one_row=true
diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs
@@ -1133,9 +1133,14 @@ pub fn project_with_column_index(
         .into_iter()
         .enumerate()
         .map(|(i, e)| match e {
+            alias @ Expr::Alias { .. }
+                if &alias.display_name().unwrap() != schema.field(i).name() =>
+            {
+                alias.unalias().alias(schema.field(i).name())
+            }
             ignore_alias @ Expr::Alias { .. } => ignore_alias,
             ignore_col @ Expr::Column { .. } => ignore_col,
-            x => x.alias(schema.field(i).name()),
+            expr => expr.alias(schema.field(i).name()),
         })
         .collect::<Vec<_>>();
     Ok(LogicalPlan::Projection(Projection::try_new_with_schema(
@@ -1187,7 +1192,7 @@ pub fn union(left_plan: LogicalPlan, right_plan: LogicalPlan) -> Result<LogicalP
         .into_iter()
         .flat_map(|p| match p {
             LogicalPlan::Union(Union { inputs, .. }) => inputs,
-            x => vec![Arc::new(x)],
+            other_plan => vec![Arc::new(other_plan)],
         })
         .map(|p| {
             let plan = coerce_plan_expr_for_schema(&p, &union_schema)?;
@@ -1199,7 +1204,7 @@ pub fn union(left_plan: LogicalPlan, right_plan: LogicalPlan) -> Result<LogicalP
                         Arc::new(union_schema.clone()),
                     )?))
                 }
-                x => Ok(Arc::new(x)),
+                other_plan => Ok(Arc::new(other_plan)),
             }
         })
         .collect::<Result<Vec<_>>>()?;