Compute variable positions for explicit projections (#2573)

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)?;

crates/physical-plan/src/plan.rs

@@ -146,12 +146,21 @@ impl ProjectListPlan {
                     .collect::<Result<_>>()?,
                 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::<Result<_>>()?,
-                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))
+            }
         }
     }