From fb1f67983d82c59bf583cc2cc810fce034b05905 Mon Sep 17 00:00:00 2001 From: Dmitrii Blaginin Date: Tue, 11 Mar 2025 13:26:05 +0000 Subject: [PATCH 1/4] Only unnest source for `EmptyRelation` --- datafusion/sql/src/unparser/plan.rs | 13 +++++++++++-- datafusion/sql/tests/cases/plan_to_sql.rs | 18 ++++++++++++++++++ 2 files changed, 29 insertions(+), 2 deletions(-) diff --git a/datafusion/sql/src/unparser/plan.rs b/datafusion/sql/src/unparser/plan.rs index b14fbdff236f5..c74b100db2eb9 100644 --- a/datafusion/sql/src/unparser/plan.rs +++ b/datafusion/sql/src/unparser/plan.rs @@ -377,8 +377,17 @@ impl Unparser<'_> { }; if self.dialect.unnest_as_table_factor() && unnest_input_type.is_some() { if let LogicalPlan::Unnest(unnest) = &p.input.as_ref() { - return self - .unnest_to_table_factor_sql(unnest, query, select, relation); + if let LogicalPlan::Projection(projection) = unnest.input.as_ref() + { + if matches!( + projection.input.as_ref(), + LogicalPlan::EmptyRelation(_) + ) { + return self.unnest_to_table_factor_sql( + unnest, query, select, relation, + ); + } + } } } diff --git a/datafusion/sql/tests/cases/plan_to_sql.rs b/datafusion/sql/tests/cases/plan_to_sql.rs index fc0b7a26baafa..aca155916c691 100644 --- a/datafusion/sql/tests/cases/plan_to_sql.rs +++ b/datafusion/sql/tests/cases/plan_to_sql.rs @@ -639,6 +639,24 @@ fn roundtrip_statement_with_dialect() -> Result<()> { parser_dialect: Box::new(GenericDialect {}), unparser_dialect: Box::new(CustomDialectBuilder::default().with_unnest_as_table_factor(true).build()), }, + TestStatementWithDialect { + sql: "SELECT unnest([1, 2, 3, 4]) from unnest([1, 2, 3]);", + expected: r#"SELECT UNNEST([1, 2, 3, 4]) AS UNNEST(make_array(Int64(1),Int64(2),Int64(3),Int64(4))) FROM UNNEST([1, 2, 3])"#, + parser_dialect: Box::new(GenericDialect {}), + unparser_dialect: Box::new(CustomDialectBuilder::default().with_unnest_as_table_factor(true).build()), + }, + TestStatementWithDialect { + sql: "SELECT unnest([1, 2, 3, 4]) from unnest([1, 2, 3]);", + expected: r#"SELECT UNNEST([1, 2, 3, 4]) AS UNNEST(make_array(Int64(1),Int64(2),Int64(3),Int64(4))) FROM UNNEST([1, 2, 3])"#, + parser_dialect: Box::new(GenericDialect {}), + unparser_dialect: Box::new(CustomDialectBuilder::default().with_unnest_as_table_factor(true).build()), + }, + TestStatementWithDialect { + sql: "select * from (select [1, 2, 3] as c), unnest(c);", + expected: r#"SELECT * FROM (SELECT [1, 2, 3] AS c) CROSS JOIN UNNEST(c)"#, + parser_dialect: Box::new(GenericDialect {}), + unparser_dialect: Box::new(CustomDialectBuilder::default().with_unnest_as_table_factor(true).build()), + }, TestStatementWithDialect { sql: "SELECT * FROM unnest_table u, UNNEST(u.array_col)", expected: r#"SELECT * FROM unnest_table AS u CROSS JOIN LATERAL (SELECT UNNEST(u.array_col) AS "UNNEST(outer_ref(u.array_col))")"#, From 28f53014482399c6f74ac2bfda464de712705663 Mon Sep 17 00:00:00 2001 From: Dmitrii Blaginin Date: Sat, 15 Mar 2025 21:07:21 +0000 Subject: [PATCH 2/4] Add a note on new condition --- datafusion/sql/src/unparser/plan.rs | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/datafusion/sql/src/unparser/plan.rs b/datafusion/sql/src/unparser/plan.rs index c74b100db2eb9..82a65a4de581a 100644 --- a/datafusion/sql/src/unparser/plan.rs +++ b/datafusion/sql/src/unparser/plan.rs @@ -381,7 +381,13 @@ impl Unparser<'_> { { if matches!( projection.input.as_ref(), - LogicalPlan::EmptyRelation(_) + LogicalPlan::EmptyRelation(_) // It may be possible that UNNEST is used as a source for the query. + // However, at this point, we don't yet know if it is just a single expression + // from another source or if it's from UNNEST. + // + // Unnest(Projection(EmptyRelation)) denotes a case with `UNNEST([...])`, + // which is normally safe to unnest as a table factor. + // However, in the future, more comprehensive checks can be added here. ) { return self.unnest_to_table_factor_sql( unnest, query, select, relation, From ef4a79dd160c452f8f92c82ac151131865c2e965 Mon Sep 17 00:00:00 2001 From: blaginin Date: Mon, 17 Mar 2025 12:11:54 +0000 Subject: [PATCH 3/4] Remove new test --- datafusion/sql/tests/cases/plan_to_sql.rs | 6 ------ 1 file changed, 6 deletions(-) diff --git a/datafusion/sql/tests/cases/plan_to_sql.rs b/datafusion/sql/tests/cases/plan_to_sql.rs index 0ab3d42a31890..25dbf51bf721c 100644 --- a/datafusion/sql/tests/cases/plan_to_sql.rs +++ b/datafusion/sql/tests/cases/plan_to_sql.rs @@ -645,12 +645,6 @@ fn roundtrip_statement_with_dialect() -> Result<()> { parser_dialect: Box::new(GenericDialect {}), unparser_dialect: Box::new(CustomDialectBuilder::default().with_unnest_as_table_factor(true).build()), }, - TestStatementWithDialect { - sql: "select * from (select [1, 2, 3] as c), unnest(c);", - expected: r#"SELECT * FROM (SELECT [1, 2, 3] AS c) CROSS JOIN UNNEST(c)"#, - parser_dialect: Box::new(GenericDialect {}), - unparser_dialect: Box::new(CustomDialectBuilder::default().with_unnest_as_table_factor(true).build()), - }, TestStatementWithDialect { sql: "SELECT * FROM unnest_table u, UNNEST(u.array_col)", expected: r#"SELECT u.array_col, u.struct_col, "UNNEST(outer_ref(u.array_col))" FROM unnest_table AS u CROSS JOIN LATERAL (SELECT UNNEST(u.array_col) AS "UNNEST(outer_ref(u.array_col))")"#, From d9218c5d177680982c6280995aeae0138602923c Mon Sep 17 00:00:00 2001 From: blaginin Date: Thu, 20 Mar 2025 09:29:13 +0000 Subject: [PATCH 4/4] Put all unnest assumptions into one function --- datafusion/sql/src/unparser/plan.rs | 53 +++++++++++++++-------------- 1 file changed, 28 insertions(+), 25 deletions(-) diff --git a/datafusion/sql/src/unparser/plan.rs b/datafusion/sql/src/unparser/plan.rs index 82a65a4de581a..20844d80332f4 100644 --- a/datafusion/sql/src/unparser/plan.rs +++ b/datafusion/sql/src/unparser/plan.rs @@ -377,22 +377,16 @@ impl Unparser<'_> { }; if self.dialect.unnest_as_table_factor() && unnest_input_type.is_some() { if let LogicalPlan::Unnest(unnest) = &p.input.as_ref() { - if let LogicalPlan::Projection(projection) = unnest.input.as_ref() + if let Some(unnest_relation) = + self.try_unnest_to_table_factor_sql(unnest)? { - if matches!( - projection.input.as_ref(), - LogicalPlan::EmptyRelation(_) // It may be possible that UNNEST is used as a source for the query. - // However, at this point, we don't yet know if it is just a single expression - // from another source or if it's from UNNEST. - // - // Unnest(Projection(EmptyRelation)) denotes a case with `UNNEST([...])`, - // which is normally safe to unnest as a table factor. - // However, in the future, more comprehensive checks can be added here. - ) { - return self.unnest_to_table_factor_sql( - unnest, query, select, relation, - ); - } + relation.unnest(unnest_relation); + return self.select_to_sql_recursively( + p.input.as_ref(), + query, + select, + relation, + ); } } } @@ -869,25 +863,34 @@ impl Unparser<'_> { None } - fn unnest_to_table_factor_sql( + fn try_unnest_to_table_factor_sql( &self, unnest: &Unnest, - query: &mut Option, - select: &mut SelectBuilder, - relation: &mut RelationBuilder, - ) -> Result<()> { + ) -> Result> { let mut unnest_relation = UnnestRelationBuilder::default(); - let LogicalPlan::Projection(p) = unnest.input.as_ref() else { - return internal_err!("Unnest input is not a Projection: {unnest:?}"); + let LogicalPlan::Projection(projection) = unnest.input.as_ref() else { + return Ok(None); + }; + + if !matches!(projection.input.as_ref(), LogicalPlan::EmptyRelation(_)) { + // It may be possible that UNNEST is used as a source for the query. + // However, at this point, we don't yet know if it is just a single expression + // from another source or if it's from UNNEST. + // + // Unnest(Projection(EmptyRelation)) denotes a case with `UNNEST([...])`, + // which is normally safe to unnest as a table factor. + // However, in the future, more comprehensive checks can be added here. + return Ok(None); }; - let exprs = p + + let exprs = projection .expr .iter() .map(|e| self.expr_to_sql(e)) .collect::>>()?; unnest_relation.array_exprs(exprs); - relation.unnest(unnest_relation); - self.select_to_sql_recursively(p.input.as_ref(), query, select, relation) + + Ok(Some(unnest_relation)) } fn is_scan_with_pushdown(scan: &TableScan) -> bool {