From 0ccc591e4c2fec372b0de7780c4f8cd9d3a57d20 Mon Sep 17 00:00:00 2001
From: comphead <comphead@ukr.net>
Date: Wed, 17 May 2023 08:59:10 -0700
Subject: [PATCH 1/3] Fix UNION ALL aliasing

---
 .../tests/sqllogictests/test_files/union.slt  | 27 +++++++++++++++++++
 datafusion/expr/src/logical_plan/builder.rs   |  3 +++
 2 files changed, 30 insertions(+)

diff --git a/datafusion/core/tests/sqllogictests/test_files/union.slt b/datafusion/core/tests/sqllogictests/test_files/union.slt
index 9015996aedf6..4c5a5c8ed843 100644
--- 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
\ No newline at end of file
diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs
index efc16fb0d693..4bd754806e0c 100644
--- a/datafusion/expr/src/logical_plan/builder.rs
+++ b/datafusion/expr/src/logical_plan/builder.rs
@@ -1133,6 +1133,7 @@ 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()),
@@ -1208,6 +1209,8 @@ pub fn union(left_plan: LogicalPlan, right_plan: LogicalPlan) -> Result<LogicalP
         return Err(DataFusionError::Plan("Empty UNION".to_string()));
     }
 
+    dbg!(&inputs);
+
     Ok(LogicalPlan::Union(Union {
         inputs,
         schema: Arc::new(union_schema),

From 882c18a13a080fd76203d6d2b263fd685ec526a8 Mon Sep 17 00:00:00 2001
From: comphead <comphead@ukr.net>
Date: Wed, 17 May 2023 09:24:26 -0700
Subject: [PATCH 2/3] fmt

---
 datafusion/expr/src/logical_plan/builder.rs | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs
index 4bd754806e0c..4fcddfaf02c6 100644
--- a/datafusion/expr/src/logical_plan/builder.rs
+++ b/datafusion/expr/src/logical_plan/builder.rs
@@ -1133,7 +1133,11 @@ 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()),
+            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()),

From 47f93a326feb12d3f4bce57aaf8ac8e26a39a527 Mon Sep 17 00:00:00 2001
From: comphead <comphead@ukr.net>
Date: Wed, 17 May 2023 09:29:16 -0700
Subject: [PATCH 3/3] fmt

---
 datafusion/expr/src/logical_plan/builder.rs | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs
index 4fcddfaf02c6..b3076778bede 100644
--- a/datafusion/expr/src/logical_plan/builder.rs
+++ b/datafusion/expr/src/logical_plan/builder.rs
@@ -1140,7 +1140,7 @@ pub fn project_with_column_index(
             }
             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(
@@ -1192,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)?;
@@ -1204,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<_>>>()?;
@@ -1213,8 +1213,6 @@ pub fn union(left_plan: LogicalPlan, right_plan: LogicalPlan) -> Result<LogicalP
         return Err(DataFusionError::Plan("Empty UNION".to_string()));
     }
 
-    dbg!(&inputs);
-
     Ok(LogicalPlan::Union(Union {
         inputs,
         schema: Arc::new(union_schema),