fix: unparse join without projection (#15693)

* fix: unparse join without projection

* update based on goldmedal's suggestion

Author: chenkovsky

datafusion/sql/src/unparser/ast.rs

@@ -155,6 +155,11 @@ impl SelectBuilder {
         self.projection = value;
         self
     }
+    pub fn pop_projections(&mut self) -> Vec<ast::SelectItem> {
+        let ret = self.projection.clone();
+        self.projection.clear();
+        ret
+    }
     pub fn already_projected(&self) -> bool {
         !self.projection.is_empty()
     }

datafusion/sql/src/unparser/plan.rs

@@ -582,6 +582,10 @@ impl Unparser<'_> {
                     }
                     _ => (&join.left, &join.right),
                 };
+                // If there's an outer projection plan, it will already set up the projection.
+                // In that case, we don't need to worry about setting up the projection here.
+                // The outer projection plan will handle projecting the correct columns.
+                let already_projected = select.already_projected();
 
                 let left_plan =
                     match try_transform_to_simple_table_scan_with_filters(left_plan)? {
@@ -599,6 +603,13 @@ impl Unparser<'_> {
                     relation,
                 )?;
 
+                let left_projection: Option<Vec<ast::SelectItem>> = if !already_projected
+                {
+                    Some(select.pop_projections())
+                } else {
+                    None
+                };
+
                 let right_plan =
                     match try_transform_to_simple_table_scan_with_filters(right_plan)? {
                         Some((plan, filters)) => {
@@ -657,6 +668,13 @@ impl Unparser<'_> {
                     &mut right_relation,
                 )?;
 
+                let right_projection: Option<Vec<ast::SelectItem>> = if !already_projected
+                {
+                    Some(select.pop_projections())
+                } else {
+                    None
+                };
+
                 match join.join_type {
                     JoinType::LeftSemi
                     | JoinType::LeftAnti
@@ -702,6 +720,9 @@ impl Unparser<'_> {
                         } else {
                             select.selection(Some(exists_expr));
                         }
+                        if let Some(projection) = left_projection {
+                            select.projection(projection);
+                        }
                     }
                     JoinType::Inner
                     | JoinType::Left
@@ -719,6 +740,21 @@ impl Unparser<'_> {
                         let mut from = select.pop_from().unwrap();
                         from.push_join(ast_join);
                         select.push_from(from);
+                        if !already_projected {
+                            let Some(left_projection) = left_projection else {
+                                return internal_err!("Left projection is missing");
+                            };
+
+                            let Some(right_projection) = right_projection else {
+                                return internal_err!("Right projection is missing");
+                            };
+
+                            let projection = left_projection
+                                .into_iter()
+                                .chain(right_projection.into_iter())
+                                .collect();
+                            select.projection(projection);
+                        }
                     }
                 };
 

datafusion/sql/src/unparser/utils.rs

@@ -385,7 +385,7 @@ pub(crate) fn try_transform_to_simple_table_scan_with_filters(
                 let mut builder = LogicalPlanBuilder::scan(
                     table_scan.table_name.clone(),
                     Arc::clone(&table_scan.source),
-                    None,
+                    table_scan.projection.clone(),
                 )?;
 
                 if let Some(alias) = table_alias.take() {

datafusion/sql/tests/cases/plan_to_sql.rs

@@ -2342,3 +2342,86 @@ fn test_unparse_right_anti_join() -> Result<()> {
     );
     Ok(())
 }
+
+#[test]
+fn test_unparse_cross_join_with_table_scan_projection() -> Result<()> {
+    let schema = Schema::new(vec![
+        Field::new("k", DataType::Int32, false),
+        Field::new("v", DataType::Int32, false),
+    ]);
+    // Cross Join:
+    //   SubqueryAlias: t1
+    //     TableScan: test projection=[v]
+    //   SubqueryAlias: t2
+    //     TableScan: test projection=[v]
+    let table_scan1 = table_scan(Some("test"), &schema, Some(vec![1]))?.build()?;
+    let table_scan2 = table_scan(Some("test"), &schema, Some(vec![1]))?.build()?;
+    let plan = LogicalPlanBuilder::from(subquery_alias(table_scan1, "t1")?)
+        .cross_join(subquery_alias(table_scan2, "t2")?)?
+        .build()?;
+    let unparser = Unparser::new(&UnparserPostgreSqlDialect {});
+    let sql = unparser.plan_to_sql(&plan)?;
+    assert_snapshot!(
+        sql,
+        @r#"SELECT "t1"."v", "t2"."v" FROM "test" AS "t1" CROSS JOIN "test" AS "t2""#
+    );
+    Ok(())
+}
+
+#[test]
+fn test_unparse_inner_join_with_table_scan_projection() -> Result<()> {
+    let schema = Schema::new(vec![
+        Field::new("k", DataType::Int32, false),
+        Field::new("v", DataType::Int32, false),
+    ]);
+    // Inner Join:
+    //   SubqueryAlias: t1
+    //     TableScan: test projection=[v]
+    //   SubqueryAlias: t2
+    //     TableScan: test projection=[v]
+    let table_scan1 = table_scan(Some("test"), &schema, Some(vec![1]))?.build()?;
+    let table_scan2 = table_scan(Some("test"), &schema, Some(vec![1]))?.build()?;
+    let plan = LogicalPlanBuilder::from(subquery_alias(table_scan1, "t1")?)
+        .join_on(
+            subquery_alias(table_scan2, "t2")?,
+            datafusion_expr::JoinType::Inner,
+            vec![col("t1.v").eq(col("t2.v"))],
+        )?
+        .build()?;
+    let unparser = Unparser::new(&UnparserPostgreSqlDialect {});
+    let sql = unparser.plan_to_sql(&plan)?;
+    assert_snapshot!(
+        sql,
+        @r#"SELECT "t1"."v", "t2"."v" FROM "test" AS "t1" INNER JOIN "test" AS "t2" ON ("t1"."v" = "t2"."v")"#
+    );
+    Ok(())
+}
+
+#[test]
+fn test_unparse_left_semi_join_with_table_scan_projection() -> Result<()> {
+    let schema = Schema::new(vec![
+        Field::new("k", DataType::Int32, false),
+        Field::new("v", DataType::Int32, false),
+    ]);
+    // LeftSemi Join:
+    //   SubqueryAlias: t1
+    //     TableScan: test projection=[v]
+    //   SubqueryAlias: t2
+    //     TableScan: test projection=[v]
+    let table_scan1 = table_scan(Some("test"), &schema, Some(vec![1]))?.build()?;
+    let table_scan2 = table_scan(Some("test"), &schema, Some(vec![1]))?.build()?;
+    let plan = LogicalPlanBuilder::from(subquery_alias(table_scan1, "t1")?)
+        .join_on(
+            subquery_alias(table_scan2, "t2")?,
+            datafusion_expr::JoinType::LeftSemi,
+            vec![col("t1.v").eq(col("t2.v"))],
+        )?
+        .build()?;
+    let unparser = Unparser::new(&UnparserPostgreSqlDialect {});
+    let sql = unparser.plan_to_sql(&plan)?;
+    assert_snapshot!(
+        sql,
+        @r#"SELECT "t1"."v" FROM "test" AS "t1" WHERE EXISTS (SELECT 1 FROM "test" AS "t2" WHERE ("t1"."v" = "t2"."v"))"#
+    );
+    Ok(())
+}