Skip to content

Commit ecb0044

Browse files
authored
Fix unnest conjunction with selecting wildcard expression (#12760)
* fix unnest statement with wildcard expression * add commnets
1 parent 9bf0630 commit ecb0044

File tree

3 files changed

+66
-2
lines changed

3 files changed

+66
-2
lines changed

datafusion/expr/src/utils.rs

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
2020
use std::cmp::Ordering;
2121
use std::collections::{HashMap, HashSet};
22+
use std::ops::Deref;
2223
use std::sync::Arc;
2324

2425
use crate::expr::{Alias, Sort, WildcardOptions, WindowFunction};
@@ -754,6 +755,15 @@ pub fn find_base_plan(input: &LogicalPlan) -> &LogicalPlan {
754755
match input {
755756
LogicalPlan::Window(window) => find_base_plan(&window.input),
756757
LogicalPlan::Aggregate(agg) => find_base_plan(&agg.input),
758+
// [SqlToRel::try_process_unnest] will convert Expr(Unnest(Expr)) to Projection/Unnest/Projection
759+
// We should expand the wildcard expression based on the input plan of the inner Projection.
760+
LogicalPlan::Unnest(unnest) => {
761+
if let LogicalPlan::Projection(projection) = unnest.input.deref() {
762+
find_base_plan(&projection.input)
763+
} else {
764+
input
765+
}
766+
}
757767
LogicalPlan::Filter(filter) => {
758768
if filter.having {
759769
// If a filter is used for a having clause, its input plan is an aggregation.

datafusion/sql/src/utils.rs

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -619,7 +619,9 @@ pub(crate) fn rewrite_recursive_unnest_bottom_up(
619619
} = original_expr.clone().rewrite(&mut rewriter)?;
620620

621621
if !transformed {
622-
if matches!(&transformed_expr, Expr::Column(_)) {
622+
if matches!(&transformed_expr, Expr::Column(_))
623+
|| matches!(&transformed_expr, Expr::Wildcard { .. })
624+
{
623625
push_projection_dedupl(inner_projection_exprs, transformed_expr.clone());
624626
Ok(vec![transformed_expr])
625627
} else {

datafusion/sqllogictest/test_files/unnest.slt

Lines changed: 53 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ AS VALUES
3333
statement ok
3434
CREATE TABLE nested_unnest_table
3535
AS VALUES
36-
(struct('a', 'b', struct('c')), (struct('a', 'b', [10,20])), [struct('a', 'b')]),
36+
(struct('a', 'b', struct('c')), (struct('a', 'b', [10,20])), [struct('a', 'b')]),
3737
(struct('d', 'e', struct('f')), (struct('x', 'y', [30,40, 50])), null)
3838
;
3939

@@ -780,3 +780,55 @@ NULL 1
780780
### TODO: group by unnest struct
781781
query error DataFusion error: Error during planning: Projection references non\-aggregate values
782782
select unnest(column1) c1 from nested_unnest_table group by c1.c0;
783+
784+
query II??I??
785+
select unnest(column5), * from unnest_table;
786+
----
787+
1 2 [1, 2, 3] [7] 1 [13, 14] {c0: 1, c1: 2}
788+
3 4 [4, 5] [8, 9, 10] 2 [15, 16] {c0: 3, c1: 4}
789+
NULL NULL [6] [11, 12] 3 NULL NULL
790+
7 8 [12] [, 42, ] NULL NULL {c0: 7, c1: 8}
791+
NULL NULL NULL NULL 4 [17, 18] NULL
792+
793+
query TT????
794+
select unnest(column1), * from nested_unnest_table
795+
----
796+
a b {c0: c} {c0: a, c1: b, c2: {c0: c}} {c0: a, c1: b, c2: [10, 20]} [{c0: a, c1: b}]
797+
d e {c0: f} {c0: d, c1: e, c2: {c0: f}} {c0: x, c1: y, c2: [30, 40, 50]} NULL
798+
799+
query ?????
800+
select unnest(unnest(column3)), * from recursive_unnest_table
801+
----
802+
[1] [[1, 2]] {c0: [1], c1: a} [[[1], [2]], [[1, 1]]] [{c0: [1], c1: [[1, 2]]}]
803+
[2] [[3], [4]] {c0: [2], c1: b} [[[3, 4], [5]], [[, 6], , [7, 8]]] [{c0: [2], c1: [[3], [4]]}]
804+
805+
statement ok
806+
CREATE TABLE join_table
807+
AS VALUES
808+
(1, 2, 3),
809+
(2, 3, 4),
810+
(4, 5, 6)
811+
;
812+
813+
query IIIII
814+
select unnest(u.column5), j.* from unnest_table u join join_table j on u.column3 = j.column1
815+
----
816+
1 2 1 2 3
817+
3 4 2 3 4
818+
NULL NULL 4 5 6
819+
820+
query II?I?
821+
select unnest(column5), * except (column5, column1) from unnest_table;
822+
----
823+
1 2 [7] 1 [13, 14]
824+
3 4 [8, 9, 10] 2 [15, 16]
825+
NULL NULL [11, 12] 3 NULL
826+
7 8 [, 42, ] NULL NULL
827+
NULL NULL NULL 4 [17, 18]
828+
829+
query III
830+
select unnest(u.column5), j.* except(column2, column3) from unnest_table u join join_table j on u.column3 = j.column1
831+
----
832+
1 2 1
833+
3 4 2
834+
NULL NULL 4

0 commit comments

Comments
 (0)