Skip to content

Commit

Permalink
fix: Add additional required expression for natural join (apache#11713)
Browse files Browse the repository at this point in the history
* adding test file

* replace expand_wildcard

* refine test

---------

Co-authored-by: jonahgao <[email protected]>
  • Loading branch information
Lordworms and jonahgao authored Aug 3, 2024
1 parent 81668f3 commit 9e90e17
Show file tree
Hide file tree
Showing 2 changed files with 54 additions and 2 deletions.
7 changes: 6 additions & 1 deletion datafusion/expr/src/logical_plan/builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1531,7 +1531,12 @@ pub fn wrap_projection_for_join_if_necessary(

let need_project = join_keys.iter().any(|key| !matches!(key, Expr::Column(_)));
let plan = if need_project {
let mut projection = expand_wildcard(input_schema, &input, None)?;
// Include all columns from the input and extend them with the join keys
let mut projection = input_schema
.columns()
.into_iter()
.map(Expr::Column)
.collect::<Vec<_>>();
let join_key_items = alias_join_keys
.iter()
.flat_map(|expr| expr.try_as_col().is_none().then_some(expr))
Expand Down
49 changes: 48 additions & 1 deletion datafusion/sqllogictest/test_files/join.slt
Original file line number Diff line number Diff line change
Expand Up @@ -1130,4 +1130,51 @@ SELECT * FROM
(SELECT * FROM t1 CROSS JOIN t2)
WHERE t1.a + t2.a IS NULL;
----
NULL NULL
NULL NULL

statement ok
CREATE TABLE t5(v0 BIGINT, v1 STRING, v2 BIGINT, v3 STRING, v4 BOOLEAN);

statement ok
CREATE TABLE t1(v0 BIGINT, v1 STRING);

statement ok
CREATE TABLE t0(v0 BIGINT, v1 DOUBLE);

query TT
explain SELECT *
FROM t1
NATURAL JOIN t5
INNER JOIN t0 ON (t0.v1 + t5.v0) > 0
WHERE t0.v1 = t1.v0;
----
logical_plan
01)Projection: t1.v0, t1.v1, t5.v2, t5.v3, t5.v4, t0.v0, t0.v1
02)--Inner Join: CAST(t1.v0 AS Float64) = t0.v1 Filter: t0.v1 + CAST(t5.v0 AS Float64) > Float64(0)
03)----Projection: t1.v0, t1.v1, t5.v0, t5.v2, t5.v3, t5.v4
04)------Inner Join: Using t1.v0 = t5.v0, t1.v1 = t5.v1
05)--------TableScan: t1 projection=[v0, v1]
06)--------TableScan: t5 projection=[v0, v1, v2, v3, v4]
07)----TableScan: t0 projection=[v0, v1]
physical_plan
01)CoalesceBatchesExec: target_batch_size=8192
02)--HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(CAST(t1.v0 AS Float64)@6, v1@1)], filter=v1@1 + CAST(v0@0 AS Float64) > 0, projection=[v0@0, v1@1, v2@3, v3@4, v4@5, v0@7, v1@8]
03)----CoalescePartitionsExec
04)------ProjectionExec: expr=[v0@0 as v0, v1@1 as v1, v0@2 as v0, v2@3 as v2, v3@4 as v3, v4@5 as v4, CAST(v0@0 AS Float64) as CAST(t1.v0 AS Float64)]
05)--------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
06)----------CoalesceBatchesExec: target_batch_size=8192
07)------------HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(v0@0, v0@0), (v1@1, v1@1)], projection=[v0@0, v1@1, v0@2, v2@4, v3@5, v4@6]
08)--------------MemoryExec: partitions=1, partition_sizes=[0]
09)--------------MemoryExec: partitions=1, partition_sizes=[0]
10)----MemoryExec: partitions=1, partition_sizes=[0]



statement ok
drop table t5;

statement ok
drop table t1;

statement ok
drop table t0;

0 comments on commit 9e90e17

Please sign in to comment.