From 75a953fbfd5fe712870509ad28990a03c116391b Mon Sep 17 00:00:00 2001 From: Sergei Grebnov Date: Fri, 11 Oct 2024 08:04:23 -0700 Subject: [PATCH] Improve UNION unparsing (#39) --- datafusion/sql/src/unparser/plan.rs | 5 +++++ datafusion/sql/tests/cases/plan_to_sql.rs | 7 +++++++ 2 files changed, 12 insertions(+) diff --git a/datafusion/sql/src/unparser/plan.rs b/datafusion/sql/src/unparser/plan.rs index 24be558148a7..54e2a1f6f17e 100644 --- a/datafusion/sql/src/unparser/plan.rs +++ b/datafusion/sql/src/unparser/plan.rs @@ -550,6 +550,11 @@ impl Unparser<'_> { ); } + // Covers cases where the UNION is a subquery and the projection is at the top level + if select.already_projected() { + return self.derive(plan, relation); + } + let input_exprs: Vec = union .inputs .iter() diff --git a/datafusion/sql/tests/cases/plan_to_sql.rs b/datafusion/sql/tests/cases/plan_to_sql.rs index fbd977875558..1eca29cfba2d 100644 --- a/datafusion/sql/tests/cases/plan_to_sql.rs +++ b/datafusion/sql/tests/cases/plan_to_sql.rs @@ -139,6 +139,13 @@ fn roundtrip_statement() -> Result<()> { SELECT j2_string as string FROM j2 ORDER BY string DESC LIMIT 10"#, + r#"SELECT col1, id FROM ( + SELECT j1_string AS col1, j1_id AS id FROM j1 + UNION ALL + SELECT j2_string AS col1, j2_id AS id FROM j2 + UNION ALL + SELECT j3_string AS col1, j3_id AS id FROM j3 + ) AS subquery GROUP BY col1, id ORDER BY col1 ASC, id ASC"#, "SELECT id, count(*) over (PARTITION BY first_name ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING), last_name, sum(id) over (PARTITION BY first_name ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING), first_name from person",