Skip to content

Commit 0536219

Browse files
authored
Fix UNION ALL aliasing (#6373)
* Fix UNION ALL aliasing * fmt * fmt
1 parent 7c6b41a commit 0536219

File tree

2 files changed

+35
-3
lines changed

2 files changed

+35
-3
lines changed

datafusion/core/tests/sqllogictests/test_files/union.slt

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -446,3 +446,30 @@ drop table t1
446446

447447
statement ok
448448
drop table t2
449+
450+
# test UNION ALL aliases correctly with all aliased
451+
query TT
452+
explain select 1 a group by a union all select 2 b union all select 3 c
453+
----
454+
logical_plan
455+
Union
456+
--Projection: Int64(1) AS a
457+
----Aggregate: groupBy=[[Int64(1)]], aggr=[[]]
458+
------EmptyRelation
459+
--Projection: Int64(2) AS a
460+
----EmptyRelation
461+
--Projection: Int64(3) AS a
462+
----EmptyRelation
463+
physical_plan
464+
UnionExec
465+
--ProjectionExec: expr=[Int64(1)@0 as a]
466+
----AggregateExec: mode=FinalPartitioned, gby=[Int64(1)@0 as Int64(1)], aggr=[]
467+
------CoalesceBatchesExec: target_batch_size=8192
468+
--------RepartitionExec: partitioning=Hash([Column { name: "Int64(1)", index: 0 }], 4), input_partitions=4
469+
----------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
470+
------------AggregateExec: mode=Partial, gby=[1 as Int64(1)], aggr=[]
471+
--------------EmptyExec: produce_one_row=true
472+
--ProjectionExec: expr=[2 as a]
473+
----EmptyExec: produce_one_row=true
474+
--ProjectionExec: expr=[3 as a]
475+
----EmptyExec: produce_one_row=true

datafusion/expr/src/logical_plan/builder.rs

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1133,9 +1133,14 @@ pub fn project_with_column_index(
11331133
.into_iter()
11341134
.enumerate()
11351135
.map(|(i, e)| match e {
1136+
alias @ Expr::Alias { .. }
1137+
if &alias.display_name().unwrap() != schema.field(i).name() =>
1138+
{
1139+
alias.unalias().alias(schema.field(i).name())
1140+
}
11361141
ignore_alias @ Expr::Alias { .. } => ignore_alias,
11371142
ignore_col @ Expr::Column { .. } => ignore_col,
1138-
x => x.alias(schema.field(i).name()),
1143+
expr => expr.alias(schema.field(i).name()),
11391144
})
11401145
.collect::<Vec<_>>();
11411146
Ok(LogicalPlan::Projection(Projection::try_new_with_schema(
@@ -1187,7 +1192,7 @@ pub fn union(left_plan: LogicalPlan, right_plan: LogicalPlan) -> Result<LogicalP
11871192
.into_iter()
11881193
.flat_map(|p| match p {
11891194
LogicalPlan::Union(Union { inputs, .. }) => inputs,
1190-
x => vec![Arc::new(x)],
1195+
other_plan => vec![Arc::new(other_plan)],
11911196
})
11921197
.map(|p| {
11931198
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
11991204
Arc::new(union_schema.clone()),
12001205
)?))
12011206
}
1202-
x => Ok(Arc::new(x)),
1207+
other_plan => Ok(Arc::new(other_plan)),
12031208
}
12041209
})
12051210
.collect::<Result<Vec<_>>>()?;

0 commit comments

Comments
 (0)