kosiew commented on code in PR #24455:
URL: https://github.com/apache/datafusion/pull/24455#discussion_r3871456881


##########
datafusion/sql/src/unparser/plan.rs:
##########
@@ -1591,10 +1591,43 @@ impl Unparser<'_> {
                     );
                 }
 
+                // Each UNION branch is unparsed in its own isolated query
+                // context. A branch is inlined only when it is a plain SELECT
+                // with no query-scoped clauses; otherwise it is wrapped in a
+                // parenthesized subquery.
+                //
+                // Sharing this statement's `QueryBuilder` across branches is
+                // unsound: its single `distinct_union` flag leaks a nested
+                // distinct `UNION` up to the enclosing `UNION ALL`, and a
+                // branch's own `ORDER BY`/`LIMIT`/`OFFSET` would bind to the
+                // whole set operation. Combined with sqlparser rendering 
nested
+                // `SetExpr::SetOperation`s without parentheses, this silently
+                // rewrites e.g. `a UNION ALL (b UNION c LIMIT 1)` into
+                // `a UNION b UNION c LIMIT 1`. Isolating each branch and
+                // parenthesizing non-trivial ones preserves precedence, the 
set
+                // quantifier, and operand-scoped clauses.
                 let input_exprs: Vec<SetExpr> = union
                     .inputs
                     .iter()
-                    .map(|input| self.select_to_sql_expr(input, query))
+                    .map(|input| {
+                        let mut branch_query = Some(QueryBuilder::default());
+                        let body = self.select_to_sql_expr(input, &mut 
branch_query)?;
+
+                        // Inline a branch only when it is a plain SELECT with 
no
+                        // query-scoped clauses; otherwise wrap it in a
+                        // parenthesized subquery so its set quantifier and
+                        // clauses stay bound to the branch.
+                        match branch_query {
+                            Some(mut branch_query)
+                                if !matches!(body, SetExpr::Select(_))

Review Comment:
   I think this wrapping condition is a little too broad. It wraps any 
non-`SetExpr::Select` body, including a nested `UNION` with the same operator 
and set quantifier as its parent. In that case the parentheses do not change 
the meaning, and they cause a regression for SQLite.
   
   For example, DataFusion plans a flat query like `a UNION ALL b UNION ALL c` 
as a nested `Union(Union(a, b), c)`. Before this change we emit the flat form, 
but with this change we emit `(a UNION ALL b) UNION ALL c`. SQLite does not 
accept a parenthesized compound SELECT in that position. I verified this with 
`SELECT 1 UNION ALL (SELECT 2 UNION ALL SELECT 3)`, which fails with a syntax 
error in SQLite.
   
   Could we wrap only when the parentheses are actually needed? Now that each 
branch has its own `QueryBuilder`, I think we can safely compute the parent's 
`set_quantifier` before the branch loop and treat a nested union with the same 
operator and quantifier as associative with the parent. Something along these 
lines:
   
   ```rust
   let associative_with_parent = matches!(
       &body,
       SetExpr::SetOperation {
           op: ast::SetOperator::Union,
           set_quantifier: sq,
           ..
       } if *sq == set_quantifier
   );
   
   let needs_wrap = !(matches!(body, SetExpr::Select(_)) || 
associative_with_parent)
       || branch_query
           .as_ref()
           .is_some_and(|q| q.has_operand_scoped_clauses());
   ```
   
   That should keep the cases this PR fixes parenthesized, such as `a UNION ALL 
(b UNION c)` and branches with their own `LIMIT` or `ORDER BY`, while allowing 
same-op/same-quantifier unions to stay flat. Could we also add a snapshot for 
the flat `a UNION ALL b UNION ALL c` case? The current round-trip tests compare 
plans, so they do not catch this SQL rendering regression.



##########
datafusion/sql/tests/cases/plan_to_sql.rs:
##########
@@ -359,6 +363,34 @@ fn roundtrip_statement_with_dialect_2() -> Result<(), 
DataFusionError> {
     Ok(())
 }
 
+#[test]
+fn roundtrip_statement_union_all_with_nested_distinct_union()
+-> Result<(), DataFusionError> {
+    // Outer `UNION ALL` whose operand is a distinct `UNION`: the outer ALL 
must
+    // survive, and the nested distinct UNION must be parenthesized.
+    roundtrip_statement_with_dialect_helper!(
+        sql: "SELECT j1_string FROM j1 UNION ALL (SELECT j2_string FROM j2 
UNION SELECT j1_string FROM j1)",
+        parser_dialect: GenericDialect {},
+        unparser_dialect: UnparserDefaultDialect {},
+        expected: @"SELECT j1.j1_string FROM j1 UNION ALL (SELECT j2.j2_string 
FROM j2 UNION SELECT j1.j1_string FROM j1)",
+    );
+    // The same shape written flat: `a UNION b UNION ALL a`.
+    roundtrip_statement_with_dialect_helper!(
+        sql: "SELECT j1_string FROM j1 UNION SELECT j2_string FROM j2 UNION 
ALL SELECT j1_string FROM j1",
+        parser_dialect: GenericDialect {},
+        unparser_dialect: UnparserDefaultDialect {},
+        expected: @"(SELECT j1.j1_string FROM j1 UNION SELECT j2.j2_string 
FROM j2) UNION ALL SELECT j1.j1_string FROM j1",
+    );
+    // A branch's own LIMIT must remain bound to that branch, inside the 
parens.
+    roundtrip_statement_with_dialect_helper!(
+        sql: "SELECT j1_string FROM j1 UNION ALL (SELECT j2_string FROM j2 
UNION SELECT j1_string FROM j1 LIMIT 5)",

Review Comment:
   Optional test hardening: we already test the branch-scoped `ORDER BY` case 
semantically through `roundtrip_statement`, but that does not pin the generated 
SQL text. It might be useful to add a 
`roundtrip_statement_with_dialect_helper!` case for something like `... UNION 
ALL (b UNION c ORDER BY 1)` so we explicitly protect the required parentheses. 
An `OFFSET` case could also cover another operand-scoped clause.



##########
datafusion/sql/src/unparser/plan.rs:
##########
@@ -1591,10 +1591,43 @@ impl Unparser<'_> {
                     );
                 }
 
+                // Each UNION branch is unparsed in its own isolated query
+                // context. A branch is inlined only when it is a plain SELECT
+                // with no query-scoped clauses; otherwise it is wrapped in a
+                // parenthesized subquery.
+                //
+                // Sharing this statement's `QueryBuilder` across branches is
+                // unsound: its single `distinct_union` flag leaks a nested
+                // distinct `UNION` up to the enclosing `UNION ALL`, and a
+                // branch's own `ORDER BY`/`LIMIT`/`OFFSET` would bind to the
+                // whole set operation. Combined with sqlparser rendering 
nested
+                // `SetExpr::SetOperation`s without parentheses, this silently
+                // rewrites e.g. `a UNION ALL (b UNION c LIMIT 1)` into
+                // `a UNION b UNION c LIMIT 1`. Isolating each branch and
+                // parenthesizing non-trivial ones preserves precedence, the 
set
+                // quantifier, and operand-scoped clauses.
                 let input_exprs: Vec<SetExpr> = union
                     .inputs
                     .iter()
-                    .map(|input| self.select_to_sql_expr(input, query))
+                    .map(|input| {
+                        let mut branch_query = Some(QueryBuilder::default());
+                        let body = self.select_to_sql_expr(input, &mut 
branch_query)?;
+
+                        // Inline a branch only when it is a plain SELECT with 
no

Review Comment:
   Minor suggestion: I think this comment repeats quite a bit of the larger 
comment just above it, as well as part of the `has_operand_scoped_clauses` doc 
comment. Could we drop this shorter comment and keep the explanation in one 
place? The part explaining why sharing the builder is unsound is especially 
useful and worth keeping.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to