[arrow-datafusion] branch main updated: Fix UNION ALL aliasing (#6373)

[alamb]

This is an automated email from the ASF dual-hosted git repository.

alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git


The following commit(s) were added to refs/heads/main by this push:
     new 0536219375 Fix UNION ALL aliasing (#6373)
0536219375 is described below

commit 0536219375444d4857e7b05a4fcfd9cdea7f480c
Author: comphead <comph...@users.noreply.github.com>
AuthorDate: Thu May 18 11:51:21 2023 -0700

    Fix UNION ALL aliasing (#6373)
    
    * Fix UNION ALL aliasing
    
    * fmt
    
    * fmt
---
 .../core/tests/sqllogictests/test_files/union.slt  | 27 ++++++++++++++++++++++
 datafusion/expr/src/logical_plan/builder.rs        | 11 ++++++---
 2 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/datafusion/core/tests/sqllogictests/test_files/union.slt 
b/datafusion/core/tests/sqllogictests/test_files/union.slt
index 9015996aed..4c5a5c8ed8 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 efc16fb0d6..b3076778be 100644
--- 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<_>>>()?;