Skip to content

Commit

Permalink
Fix unnest conjunction with selecting wildcard expression (#12760)
Browse files Browse the repository at this point in the history
* fix unnest statement with wildcard expression

* add commnets
  • Loading branch information
goldmedal authored Oct 6, 2024
1 parent 9bf0630 commit ecb0044
Show file tree
Hide file tree
Showing 3 changed files with 66 additions and 2 deletions.
10 changes: 10 additions & 0 deletions datafusion/expr/src/utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@

use std::cmp::Ordering;
use std::collections::{HashMap, HashSet};
use std::ops::Deref;
use std::sync::Arc;

use crate::expr::{Alias, Sort, WildcardOptions, WindowFunction};
Expand Down Expand Up @@ -754,6 +755,15 @@ pub fn find_base_plan(input: &LogicalPlan) -> &LogicalPlan {
match input {
LogicalPlan::Window(window) => find_base_plan(&window.input),
LogicalPlan::Aggregate(agg) => find_base_plan(&agg.input),
// [SqlToRel::try_process_unnest] will convert Expr(Unnest(Expr)) to Projection/Unnest/Projection
// We should expand the wildcard expression based on the input plan of the inner Projection.
LogicalPlan::Unnest(unnest) => {
if let LogicalPlan::Projection(projection) = unnest.input.deref() {
find_base_plan(&projection.input)
} else {
input
}
}
LogicalPlan::Filter(filter) => {
if filter.having {
// If a filter is used for a having clause, its input plan is an aggregation.
Expand Down
4 changes: 3 additions & 1 deletion datafusion/sql/src/utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -619,7 +619,9 @@ pub(crate) fn rewrite_recursive_unnest_bottom_up(
} = original_expr.clone().rewrite(&mut rewriter)?;

if !transformed {
if matches!(&transformed_expr, Expr::Column(_)) {
if matches!(&transformed_expr, Expr::Column(_))
|| matches!(&transformed_expr, Expr::Wildcard { .. })
{
push_projection_dedupl(inner_projection_exprs, transformed_expr.clone());
Ok(vec![transformed_expr])
} else {
Expand Down
54 changes: 53 additions & 1 deletion datafusion/sqllogictest/test_files/unnest.slt
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ AS VALUES
statement ok
CREATE TABLE nested_unnest_table
AS VALUES
(struct('a', 'b', struct('c')), (struct('a', 'b', [10,20])), [struct('a', 'b')]),
(struct('a', 'b', struct('c')), (struct('a', 'b', [10,20])), [struct('a', 'b')]),
(struct('d', 'e', struct('f')), (struct('x', 'y', [30,40, 50])), null)
;

Expand Down Expand Up @@ -780,3 +780,55 @@ NULL 1
### TODO: group by unnest struct
query error DataFusion error: Error during planning: Projection references non\-aggregate values
select unnest(column1) c1 from nested_unnest_table group by c1.c0;

query II??I??
select unnest(column5), * from unnest_table;
----
1 2 [1, 2, 3] [7] 1 [13, 14] {c0: 1, c1: 2}
3 4 [4, 5] [8, 9, 10] 2 [15, 16] {c0: 3, c1: 4}
NULL NULL [6] [11, 12] 3 NULL NULL
7 8 [12] [, 42, ] NULL NULL {c0: 7, c1: 8}
NULL NULL NULL NULL 4 [17, 18] NULL

query TT????
select unnest(column1), * from nested_unnest_table
----
a b {c0: c} {c0: a, c1: b, c2: {c0: c}} {c0: a, c1: b, c2: [10, 20]} [{c0: a, c1: b}]
d e {c0: f} {c0: d, c1: e, c2: {c0: f}} {c0: x, c1: y, c2: [30, 40, 50]} NULL

query ?????
select unnest(unnest(column3)), * from recursive_unnest_table
----
[1] [[1, 2]] {c0: [1], c1: a} [[[1], [2]], [[1, 1]]] [{c0: [1], c1: [[1, 2]]}]
[2] [[3], [4]] {c0: [2], c1: b} [[[3, 4], [5]], [[, 6], , [7, 8]]] [{c0: [2], c1: [[3], [4]]}]

statement ok
CREATE TABLE join_table
AS VALUES
(1, 2, 3),
(2, 3, 4),
(4, 5, 6)
;

query IIIII
select unnest(u.column5), j.* from unnest_table u join join_table j on u.column3 = j.column1
----
1 2 1 2 3
3 4 2 3 4
NULL NULL 4 5 6

query II?I?
select unnest(column5), * except (column5, column1) from unnest_table;
----
1 2 [7] 1 [13, 14]
3 4 [8, 9, 10] 2 [15, 16]
NULL NULL [11, 12] 3 NULL
7 8 [, 42, ] NULL NULL
NULL NULL NULL 4 [17, 18]

query III
select unnest(u.column5), j.* except(column2, column3) from unnest_table u join join_table j on u.column3 = j.column1
----
1 2 1
3 4 2
NULL NULL 4

0 comments on commit ecb0044

Please sign in to comment.