Skip to content

Fix unparse union all nested distinct #24455

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Jeadie wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Fix unparse union all nested distinct #24455

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
cdf866f
fix: unparser drops UNION ALL when its operand is a distinct UNION
Jeadie Aug 18, 2026
755bde7
generalize: isolate every UNION branch when unparsing
Jeadie Aug 18, 2026
76e8907
refactor: express UNION branch wrapping without expect/double-negation
Jeadie Aug 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 15 additions & 0 deletions datafusion/sql/src/unparser/ast.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,21 @@ impl QueryBuilder {
pub fn is_distinct_union(&self) -> bool {
self.distinct_union
}
/// Whether this query carries clauses that are scoped to a single query
/// rather than to a set operation. When such a query is used as an operand
/// of a set operation (e.g. a branch of a `UNION`), it must be wrapped in
/// parentheses so these clauses bind to the operand and not to the whole
/// set operation.
pub fn has_operand_scoped_clauses(&self) -> bool {
self.with.is_some()
|| self.order_by_kind.is_some()
|| self.limit.is_some()
|| self.offset.is_some()
|| self.fetch.is_some()
|| self.for_clause.is_some()
|| !self.limit_by.is_empty()
|| !self.locks.is_empty()
}
pub fn build(&self) -> Result<ast::Query, BuilderError> {
let order_by = self
.order_by_kind
Expand Down
35 changes: 34 additions & 1 deletion datafusion/sql/src/unparser/plan.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

@kosiew kosiew Aug 27, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

// 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(_))

@kosiew kosiew Aug 27, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

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.

|| branch_query.has_operand_scoped_clauses() =>
{
let query = branch_query.body(Box::new(body)).build()?;
Ok(SetExpr::Query(Box::new(query)))
}
_ => Ok(body),
}
})
.collect::<Result<Vec<_>>>()?;

assert_or_internal_err!(
Expand Down
32 changes: 32 additions & 0 deletions datafusion/sql/tests/cases/plan_to_sql.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -170,6 +170,10 @@ fn roundtrip_statement() -> Result<()> {
SELECT j2_string as string FROM j2
ORDER BY string DESC
LIMIT 10"#,
r#"SELECT j1_string FROM j1 UNION ALL (SELECT j2_string FROM j2 UNION SELECT j1_string FROM j1)"#,
r#"SELECT j1_string FROM j1 UNION SELECT j2_string FROM j2 UNION ALL SELECT j1_string FROM j1"#,
r#"SELECT j1_string FROM j1 UNION ALL (SELECT j2_string FROM j2 UNION SELECT j1_string FROM j1 LIMIT 5)"#,
r#"SELECT j1_string FROM j1 UNION ALL (SELECT j2_string FROM j2 UNION SELECT j1_string FROM j1 ORDER BY 1)"#,
r#"SELECT col1, id FROM (
SELECT j1_string AS col1, j1_id AS id FROM j1
UNION ALL
Expand Down Expand Up @@ -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)",

@kosiew kosiew Aug 27, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

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 LIMIT 5)",
);
Ok(())
}

#[test]
fn roundtrip_statement_with_dialect_3() -> Result<(), DataFusionError> {
roundtrip_statement_with_dialect_helper!(
Expand Down