fix: Add additional required expression for natural join (#11713)

* adding test file

* replace expand_wildcard

* refine test

---------

Co-authored-by: jonahgao <[email protected]>

Author: Lordworms

datafusion/expr/src/logical_plan/builder.rs

@@ -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))

datafusion/sqllogictest/test_files/join.slt

@@ -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;