diff --git a/crates/core/src/sql/execute.rs b/crates/core/src/sql/execute.rs index 652ff461fc..a2f4199c86 100644 --- a/crates/core/src/sql/execute.rs +++ b/crates/core/src/sql/execute.rs @@ -684,6 +684,33 @@ pub(crate) mod tests { Ok(()) } + /// Test projecting columns from both tables in join + #[test] + fn test_project_join() -> anyhow::Result<()> { + let db = TestDB::in_memory()?; + + let t_schema = [("id", AlgebraicType::U8), ("x", AlgebraicType::U8)]; + let s_schema = [("id", AlgebraicType::U8), ("y", AlgebraicType::U8)]; + + let t_id = db.create_table_for_test("t", &t_schema, &[0.into()])?; + let s_id = db.create_table_for_test("s", &s_schema, &[0.into()])?; + + insert_rows(&db, t_id, [product![1_u8, 2_u8]])?; + insert_rows(&db, s_id, [product![1_u8, 3_u8]])?; + + let id = identity_from_u8(1); + let auth = AuthCtx::new(Identity::ZERO, id); + + assert_query_results( + &db, + "select t.x, s.y from t join s on t.id = s.id", + &auth, + [product![2_u8, 3_u8]], + ); + + Ok(()) + } + #[test] fn test_select_star_table() -> ResultTest<()> { let (db, input) = create_data(1)?; diff --git a/crates/physical-plan/src/plan.rs b/crates/physical-plan/src/plan.rs index a505baef4a..5df081ed03 100644 --- a/crates/physical-plan/src/plan.rs +++ b/crates/physical-plan/src/plan.rs @@ -146,12 +146,21 @@ impl ProjectListPlan { .collect::>()?, agg_type, )), - Self::List(plan, fields) => Ok(Self::List( - plan.into_iter() - .map(|plan| plan.optimize(fields.iter().map(|TupleField { label, .. }| label).copied().collect())) - .collect::>()?, - fields, - )), + Self::List(plans, mut fields) => { + let mut optimized_plans = Vec::with_capacity(plans.len()); + for plan in plans { + // Collect the names of the relvars + let labels = fields.iter().map(|field| field.label).collect(); + // Optimize each plan + let optimized_plan = plan.optimize(labels)?; + // Compute the position of each relvar referenced in the projection + for TupleField { label, label_pos, .. } in &mut fields { + *label_pos = optimized_plan.position(label); + } + optimized_plans.push(optimized_plan); + } + Ok(Self::List(optimized_plans, fields)) + } } }