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(())
+}