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]