From 7f9c029f53053e802710f740a6a47acae415375c Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Fri, 2 Oct 2026 16:54:27 +0100 Subject: [PATCH 1/2] fix: keep set operations after subquery operands at the outer query A parenthesized query used as an expression operand, such as the scalar subquery in `WHERE x < (SELECT ...) UNION ALL SELECT ...`, absorbed the following set operation, so the UNION/EXCEPT/INTERSECT became part of the predicate instead of combining the outer SELECTs. The same happened in HAVING, JOIN ON, BETWEEN and doubly parenthesized operands. A parenthesized query now continues into a set operation only where it starts the contents of an enclosing parenthesis, as in `x IN ((SELECT 1) UNION (SELECT 2))`, or where CREATE TABLE AS parses its parenthesized query. Statement-level parenthesized set operations are unchanged. --- crates/polyglot-sql/src/parser.rs | 39 +++- crates/polyglot-sql/src/validation/tests.rs | 48 +++++ ...t_operation_subquery_operand_regression.rs | 185 ++++++++++++++++++ 3 files changed, 263 insertions(+), 9 deletions(-) create mode 100644 crates/polyglot-sql/tests/set_operation_subquery_operand_regression.rs diff --git a/crates/polyglot-sql/src/parser.rs b/crates/polyglot-sql/src/parser.rs index 8142778d..639a21fb 100644 --- a/crates/polyglot-sql/src/parser.rs +++ b/crates/polyglot-sql/src/parser.rs @@ -13645,9 +13645,11 @@ impl Parser { // - (WITH cte AS (SELECT 1) SELECT * FROM cte) (CTE in parens) let mut as_select_parenthesized = self.check(TokenType::LParen); let query = if as_select_parenthesized { - // Parenthesized query - parse as expression which handles subqueries - // Note: parse_primary will consume set operations like UNION internally - let subquery = self.parse_primary()?; + // Parenthesized query - parse as expression which handles subqueries and + // consumes set operations like `(SELECT 1) UNION ALL (SELECT 2)` + self.expect(TokenType::LParen)?; + let lparen_comments = self.previous_trailing_comments().to_vec(); + let subquery = self.parse_parenthesized_primary(lparen_comments, true)?; // If parse_primary returned a set operation, the outer parens weren't wrapping // the entire expression - they were part of the operands if matches!( @@ -31954,7 +31956,17 @@ impl Parser { } } - fn parse_parenthesized_primary(&mut self, lparen_comments: Vec) -> Result { + /// Parse the remainder of a parenthesized primary after its opening parenthesis. + /// + /// `set_operation_may_follow` controls whether a parenthesized query may absorb a + /// following UNION/INTERSECT/EXCEPT. Callers enable it only where the parenthesized + /// query starts a query expression; as an operator or clause operand + /// (`WHERE x < (SELECT ...) UNION ...`) the set operation belongs to the enclosing query. + fn parse_parenthesized_primary( + &mut self, + lparen_comments: Vec, + set_operation_may_follow: bool, + ) -> Result { let raw_start = self.previous().span.start; // Empty parens () — could be empty tuple or zero-param lambda () -> body @@ -32155,7 +32167,11 @@ impl Parser { })) }; - let set_result = self.parse_set_operation(subquery)?; + let set_result = if set_operation_may_follow { + self.parse_set_operation(subquery)? + } else { + subquery + }; let had_set_operation = matches!( &set_result, Expression::Union(_) | Expression::Intersect(_) | Expression::Except(_) @@ -32252,9 +32268,10 @@ impl Parser { self.expect(TokenType::RParen)?; let mut nested_paren_comments = lparen_comments.clone(); nested_paren_comments.extend_from_slice(self.previous_trailing_comments()); - if self.check(TokenType::Union) - || self.check(TokenType::Intersect) - || self.check(TokenType::Except) + if set_operation_may_follow + && (self.check(TokenType::Union) + || self.check(TokenType::Intersect) + || self.check(TokenType::Except)) { if let Expression::Subquery(subq) = &result { let set_result = self.parse_set_operation(subq.this.clone())?; @@ -33282,7 +33299,11 @@ impl Parser { // Parenthesized expression or subquery if self.match_token(TokenType::LParen) { let lparen_comments = self.previous_trailing_comments().to_vec(); - return self.parse_parenthesized_primary(lparen_comments); + // A parenthesized query that opens an enclosing parenthesis, as in + // `x IN ((SELECT 1) UNION (SELECT 2))`, may continue into a set operation. + let set_operation_may_follow = + self.current < 2 || self.tokens[self.current - 2].token_type == TokenType::LParen; + return self.parse_parenthesized_primary(lparen_comments, set_operation_may_follow); } // NULL diff --git a/crates/polyglot-sql/src/validation/tests.rs b/crates/polyglot-sql/src/validation/tests.rs index 7a344dc1..172761ff 100644 --- a/crates/polyglot-sql/src/validation/tests.rs +++ b/crates/polyglot-sql/src/validation/tests.rs @@ -4481,3 +4481,51 @@ fn schema_validation_propagates_complexity_guard_to_both_parses() { assert!(invalid.errors.iter().any(|error| error.code == "E201")); } } + +#[test] +fn set_operation_after_subquery_operand_checks_outer_branch_arity() { + let schema = semantic_type_schema(); + let options = SchemaValidationOptions { + check_types: true, + ..Default::default() + }; + for dialect in [ + DialectType::DuckDB, + DialectType::PostgreSQL, + DialectType::Snowflake, + DialectType::BigQuery, + DialectType::TSQL, + DialectType::Databricks, + ] { + for sql in [ + "SELECT o.i, o.s FROM orders AS o WHERE o.i < (SELECT MAX(p.i) FROM orders AS p) UNION ALL SELECT o.i, o.s FROM orders AS o", + "SELECT o.i, o.s FROM orders AS o WHERE o.i < (SELECT MAX(p.i) FROM orders AS p) EXCEPT SELECT o.i, o.s FROM orders AS o", + "SELECT o.i, o.s FROM orders AS o WHERE o.i < (SELECT MAX(p.i) FROM orders AS p) INTERSECT SELECT o.i, o.s FROM orders AS o", + "SELECT o.i, SUM(o.i) AS total FROM orders AS o GROUP BY o.i HAVING SUM(o.i) > (SELECT MIN(p.i) FROM orders AS p) UNION ALL SELECT o.i, o.i FROM orders AS o", + ] { + let result = validate_with_schema(sql, dialect, &schema, &options); + assert!( + !result + .errors + .iter() + .any(|error| error.message.contains("different column counts")), + "{dialect:?}: {sql}: {:?}", + result.errors + ); + } + let mismatched = validate_with_schema( + "SELECT o.i, o.s FROM orders AS o WHERE o.i < (SELECT MAX(p.i) FROM orders AS p) UNION ALL SELECT o.i FROM orders AS o", + dialect, + &schema, + &options, + ); + assert!( + mismatched + .errors + .iter() + .any(|error| error.message.contains("left 2, right 1")), + "{dialect:?}: {:?}", + mismatched.errors + ); + } +} diff --git a/crates/polyglot-sql/tests/set_operation_subquery_operand_regression.rs b/crates/polyglot-sql/tests/set_operation_subquery_operand_regression.rs new file mode 100644 index 00000000..fffb9003 --- /dev/null +++ b/crates/polyglot-sql/tests/set_operation_subquery_operand_regression.rs @@ -0,0 +1,185 @@ +//! A parenthesized query used as an expression operand must not absorb a set operation +//! that follows it. The set operation belongs to the enclosing query: +//! `SELECT ... WHERE x < (SELECT ...) UNION ALL SELECT ...` is a UNION of two SELECTs. + +use polyglot_sql::traversal::ExpressionWalk; +use polyglot_sql::{format, generate, parse_one, DialectType, Expression}; + +const DIALECTS: [DialectType; 7] = [ + DialectType::Generic, + DialectType::DuckDB, + DialectType::Snowflake, + DialectType::BigQuery, + DialectType::PostgreSQL, + DialectType::TSQL, + DialectType::Databricks, +]; + +fn is_set_operation(expression: &Expression) -> bool { + matches!( + expression, + Expression::Union(_) | Expression::Intersect(_) | Expression::Except(_) + ) +} + +fn set_operation_count(expression: &Expression) -> usize { + expression + .dfs() + .filter(|node| is_set_operation(node)) + .count() +} + +fn set_operation_operands(expression: &Expression) -> (&Expression, &Expression) { + match expression { + Expression::Union(union) => (&union.left, &union.right), + Expression::Intersect(intersect) => (&intersect.left, &intersect.right), + Expression::Except(except) => (&except.left, &except.right), + other => panic!("expected a top-level set operation, got {other:?}"), + } +} + +fn parse_stable(sql: &str, dialect: DialectType) -> Expression { + let parsed = + parse_one(sql, dialect).unwrap_or_else(|error| panic!("{dialect:?}: {sql}: {error}")); + let generated = generate(&parsed, dialect).unwrap(); + let reparsed = parse_one(&generated, dialect) + .unwrap_or_else(|error| panic!("{dialect:?}: {generated}: {error}")); + assert_eq!( + generate(&reparsed, dialect).unwrap(), + generated, + "unstable generation: {dialect:?}: {sql}" + ); + let formatted = format(&generated, dialect).unwrap().remove(0); + let formatted_ast = parse_one(&formatted, dialect) + .unwrap_or_else(|error| panic!("{dialect:?}: {formatted}: {error}")); + assert_eq!( + generate(&formatted_ast, dialect).unwrap(), + generated, + "formatter changed the query: {dialect:?}: {sql}" + ); + parsed +} + +#[test] +fn set_operation_after_subquery_operand_applies_to_enclosing_query() { + let cases = [ + "SELECT o.id, o.total FROM orders AS o WHERE o.total < (SELECT MAX(l.amount) FROM limits AS l) UNION ALL SELECT o.id, o.total FROM orders AS o", + "SELECT o.id FROM orders AS o WHERE o.total < (SELECT MAX(l.amount) FROM limits AS l) UNION SELECT o.id FROM orders AS o", + "SELECT o.id FROM orders AS o WHERE o.total < (SELECT MAX(l.amount) FROM limits AS l) EXCEPT SELECT o.id FROM orders AS o", + "SELECT o.id FROM orders AS o WHERE o.total < (SELECT MAX(l.amount) FROM limits AS l) INTERSECT SELECT o.id FROM orders AS o", + "SELECT o.id FROM orders AS o WHERE o.total < (SELECT l.amount FROM limits AS l LIMIT 1) UNION ALL SELECT o.id FROM orders AS o", + "SELECT o.id FROM orders AS o WHERE o.total < ((SELECT MAX(l.amount) FROM limits AS l)) UNION ALL SELECT o.id FROM orders AS o", + "SELECT o.id FROM orders AS o WHERE o.total < (SELECT MAX(l.amount) FROM limits AS l) AND o.qty > (SELECT 2) UNION ALL SELECT o.id FROM orders AS o", + "SELECT o.id FROM orders AS o WHERE o.total BETWEEN (SELECT 1) AND (SELECT 2) UNION ALL SELECT o.id FROM orders AS o", + "SELECT o.id FROM orders AS o WHERE o.total = (SELECT MAX(l.amount) FROM limits AS l WHERE l.amount < (SELECT 5)) UNION ALL SELECT o.id FROM orders AS o", + "SELECT o.id FROM orders AS o WHERE o.id IN (SELECT c.order_id FROM customers AS c) UNION ALL SELECT o.id FROM orders AS o", + "SELECT o.id FROM orders AS o WHERE EXISTS (SELECT 1 FROM customers AS c) UNION ALL SELECT o.id FROM orders AS o", + "SELECT o.id FROM orders AS o WHERE NOT EXISTS (SELECT 1 FROM customers AS c) EXCEPT SELECT o.id FROM orders AS o", + "SELECT (SELECT MAX(l.amount) FROM limits AS l) AS cap UNION ALL SELECT 2", + "SELECT o.id, (SELECT MAX(l.amount) FROM limits AS l) + 1 AS cap FROM orders AS o UNION ALL SELECT 1, 2", + "SELECT o.id FROM orders AS o GROUP BY o.id HAVING COUNT(*) > (SELECT 1) UNION ALL SELECT o.id FROM orders AS o", + "SELECT o.id FROM orders AS o JOIN customers AS c ON o.customer_id = (SELECT MIN(c2.id) FROM customers AS c2) UNION ALL SELECT o.id FROM orders AS o", + "SELECT o.id FROM orders AS o WHERE (SELECT TRUE) UNION ALL SELECT o.id FROM orders AS o", + ]; + for dialect in DIALECTS { + for sql in cases { + let parsed = parse_stable(sql, dialect); + let (left, right) = set_operation_operands(&parsed); + assert!( + matches!(left, Expression::Select(_)) && matches!(right, Expression::Select(_)), + "{dialect:?}: {sql}: {parsed:?}" + ); + assert_eq!( + set_operation_count(&parsed), + 1, + "{dialect:?}: set operation absorbed by a subquery operand: {sql}" + ); + } + } +} + +#[test] +fn set_operation_inside_enclosing_parentheses_stays_nested() { + let cases = [ + "SELECT o.id FROM orders AS o WHERE o.id IN ((SELECT 1) UNION ALL (SELECT 2))", + "SELECT o.id FROM orders AS o WHERE o.id IN ((SELECT 1) UNION ALL SELECT 2)", + "SELECT o.id FROM orders AS o WHERE o.id IN (SELECT 1 UNION ALL SELECT 2)", + "SELECT o.id FROM orders AS o WHERE o.id = ((SELECT 1) EXCEPT (SELECT 2))", + "SELECT o.id FROM orders AS o WHERE EXISTS ((SELECT 1) INTERSECT (SELECT 2))", + "SELECT * FROM ((SELECT 1 AS id) UNION ALL (SELECT 2 AS id)) AS t", + "WITH t AS ((SELECT 1 AS id) UNION ALL (SELECT 2 AS id)) SELECT * FROM t", + ]; + for dialect in DIALECTS { + for sql in cases { + let parsed = parse_stable(sql, dialect); + assert!( + matches!(parsed, Expression::Select(_)), + "{dialect:?}: {sql}: {parsed:?}" + ); + assert_eq!(set_operation_count(&parsed), 1, "{dialect:?}: {sql}"); + } + } +} + +#[test] +fn nested_set_operation_operand_and_outer_set_operation_keep_their_levels() { + let sql = "SELECT o.id FROM orders AS o WHERE o.id < ((SELECT 1) UNION ALL (SELECT 2)) UNION ALL SELECT 3"; + for dialect in DIALECTS { + let parsed = parse_stable(sql, dialect); + let (left, right) = set_operation_operands(&parsed); + assert!(matches!(right, Expression::Select(_)), "{dialect:?}"); + let Expression::Select(select) = left else { + panic!("{dialect:?}: expected SELECT on the left, got {left:?}") + }; + let predicate = &select.where_clause.as_ref().unwrap().this; + assert_eq!(set_operation_count(predicate), 1, "{dialect:?}"); + } +} + +#[test] +fn parenthesized_query_set_operations_keep_their_meaning() { + let cases = [ + ("(SELECT 1) UNION ALL (SELECT 2)", 1), + ("(SELECT 1) UNION ALL SELECT 2", 1), + ("SELECT 1 UNION ALL (SELECT 2)", 1), + ("(SELECT 1) EXCEPT (SELECT 2)", 1), + ("(SELECT 1) INTERSECT (SELECT 2)", 1), + ("(SELECT 1) UNION ALL (SELECT 2) UNION ALL (SELECT 3)", 2), + ("((SELECT 1) UNION ALL (SELECT 2)) UNION ALL (SELECT 3)", 2), + ]; + for dialect in DIALECTS { + for (sql, expected_set_operations) in cases { + let parsed = parse_stable(sql, dialect); + assert!(is_set_operation(&parsed), "{dialect:?}: {sql}: {parsed:?}"); + assert_eq!( + set_operation_count(&parsed), + expected_set_operations, + "{dialect:?}: {sql}" + ); + } + } +} + +#[test] +fn create_table_as_parenthesized_set_operation_keeps_set_operation() { + let sql = "CREATE TABLE t AS (SELECT 1 AS id) UNION ALL (SELECT 2 AS id)"; + for dialect in [ + DialectType::Generic, + DialectType::DuckDB, + DialectType::Snowflake, + DialectType::PostgreSQL, + DialectType::Databricks, + ] { + let parsed = parse_one(sql, dialect).unwrap_or_else(|error| panic!("{dialect:?}: {error}")); + let Expression::CreateTable(create) = &parsed else { + panic!("{dialect:?}: expected CREATE TABLE, got {parsed:?}") + }; + let query = create.as_select.as_ref().unwrap(); + assert!(is_set_operation(query), "{dialect:?}: {query:?}"); + assert_eq!( + generate(&parsed, dialect).unwrap(), + sql, + "{dialect:?}: CREATE TABLE AS changed" + ); + } +} From 9ba2a98e1a9df2be7b33d105a0e5e4e15294632a Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Fri, 2 Oct 2026 19:57:13 +0100 Subject: [PATCH 2/2] perf: build set-operation scopes in linear time Each set-operation branch scope copied its whole operand, so a chain of N UNION/INTERSECT/EXCEPT branches copied every earlier branch at each level and scope building, validation and query analysis grew quadratically. A branch that is itself an unmodified set operation now stores its operator and clauses with NULL placeholders for set-operation operands, which are the expressions of its own branch scopes. Lineage and query-analysis nullability borrow each branch's complete query from its parent instead, and Scope::output_columns follows the leftmost branch scope. The scope tree shape and all validation, analysis and lineage results are unchanged. Add set_operation_branch_queries for consumers that need a branch's complete query. --- crates/polyglot-sql/src/lib.rs | 4 +- crates/polyglot-sql/src/lineage.rs | 52 ++++-- crates/polyglot-sql/src/query_analysis.rs | 59 ++++-- crates/polyglot-sql/src/scope.rs | 193 ++++++++++++++++++-- crates/polyglot-sql/src/validation/tests.rs | 109 +++++++++++ 5 files changed, 378 insertions(+), 39 deletions(-) diff --git a/crates/polyglot-sql/src/lib.rs b/crates/polyglot-sql/src/lib.rs index c81f8a9c..74874272 100644 --- a/crates/polyglot-sql/src/lib.rs +++ b/crates/polyglot-sql/src/lib.rs @@ -133,8 +133,8 @@ pub use schema::{ }; #[cfg(feature = "semantic")] pub use scope::{ - build_scope, find_all_in_scope, find_in_scope, traverse_scope, walk_in_scope, ColumnRef, Scope, - ScopeType, SourceInfo, + build_scope, find_all_in_scope, find_in_scope, set_operation_branch_queries, traverse_scope, + walk_in_scope, ColumnRef, Scope, ScopeType, SourceInfo, }; #[cfg(feature = "time")] pub use time::{format_time, is_valid_timezone, subsecond_precision, TIMEZONES}; diff --git a/crates/polyglot-sql/src/lineage.rs b/crates/polyglot-sql/src/lineage.rs index 8cca2a70..e805bac9 100644 --- a/crates/polyglot-sql/src/lineage.rs +++ b/crates/polyglot-sql/src/lineage.rs @@ -1036,6 +1036,9 @@ struct ScopeId(usize); struct IndexedScope<'a> { scope: &'a Scope, + /// The scope's complete query. Set-operation branches borrow it from their + /// parent because nested branch scopes store a shallow copy. + expression: &'a Expression, subquery_scopes: Vec, derived_table_scopes: Vec, cte_scopes: Vec, @@ -1104,9 +1107,21 @@ impl<'a> ScopedLineage<'a> { inherited_ctes: &[&'a Scope], dialect: DialectType, schema: Option<&'a dyn Schema>, + ) -> Self { + Self::with_schema_for_expression(scope, &scope.expression, inherited_ctes, dialect, schema) + } + + /// Like [`Self::with_schema`], for a scope whose complete query is `expression`; + /// a nested set-operation branch scope stores only a shallow copy. + pub(crate) fn with_schema_for_expression( + scope: &'a Scope, + expression: &'a Expression, + inherited_ctes: &[&'a Scope], + dialect: DialectType, + schema: Option<&'a dyn Schema>, ) -> Self { let visible_ctes = scope.cte_sources.clone(); - let (mut context, root) = LineageScopeContext::from_scope(scope); + let (mut context, root) = LineageScopeContext::from_scope(scope, expression); context.schema = schema; context.conservative = true; let mut ctes = context.indexed(root).cte_scopes.clone(); @@ -1119,7 +1134,7 @@ impl<'a> ScopedLineage<'a> { .iter() .any(|id| context.scope(*id).expression == cte_scope.expression) { - ctes.push(context.insert_scope(cte_scope)); + ctes.push(context.insert_scope(cte_scope, &cte_scope.expression)); } } } @@ -1196,7 +1211,7 @@ impl<'a> ScopedLineage<'a> { pub(crate) fn output_names(&self) -> Vec { crate::ast_transforms::get_output_column_names_for_dialect( - effective_scope_expression(&self.context.scope(self.root).expression), + effective_scope_expression(self.context.expression(self.root)), self.dialect, ) } @@ -1225,7 +1240,7 @@ impl<'a> ScopedLineage<'a> { } impl<'a> LineageScopeContext<'a> { - fn from_scope(scope: &'a Scope) -> (Self, ScopeId) { + fn from_scope(scope: &'a Scope, expression: &'a Expression) -> (Self, ScopeId) { let mut context = Self { projection_indexes: RefCell::new(HashMap::new()), scopes: Vec::new(), @@ -1235,11 +1250,11 @@ impl<'a> LineageScopeContext<'a> { conservative: false, schema: None, }; - let root = context.insert_scope(scope); + let root = context.insert_scope(scope, expression); (context, root) } - fn insert_scope(&mut self, scope: &'a Scope) -> ScopeId { + fn insert_scope(&mut self, scope: &'a Scope, expression: &'a Expression) -> ScopeId { let scope_key = scope as *const Scope; if let Some(id) = self.scope_ids.get(&scope_key) { return *id; @@ -1248,6 +1263,7 @@ impl<'a> LineageScopeContext<'a> { self.scope_ids.insert(scope_key, id); self.scopes.push(IndexedScope { scope, + expression, subquery_scopes: Vec::new(), derived_table_scopes: Vec::new(), cte_scopes: Vec::new(), @@ -1256,26 +1272,32 @@ impl<'a> LineageScopeContext<'a> { let subquery_scopes = scope .subquery_scopes .iter() - .map(|child| self.insert_scope(child)) + .map(|child| self.insert_scope(child, &child.expression)) .collect(); let derived_table_scopes = scope .derived_table_scopes .iter() - .map(|child| self.insert_scope(child)) + .map(|child| self.insert_scope(child, &child.expression)) .collect(); let cte_scopes = scope .cte_scopes .iter() - .map(|child| self.insert_scope(child)) + .map(|child| self.insert_scope(child, &child.expression)) .collect(); + let branch_queries = crate::scope::set_operation_branch_queries(expression); let union_scopes = scope .union_scopes .iter() - .map(|child| self.insert_scope(child)) + .enumerate() + .map(|(index, child)| { + let query = branch_queries.map_or(&child.expression, |queries| queries[index]); + self.insert_scope(child, query) + }) .collect(); self.scopes[id.0] = IndexedScope { scope, + expression, subquery_scopes, derived_table_scopes, cte_scopes, @@ -1292,6 +1314,10 @@ impl<'a> LineageScopeContext<'a> { self.indexed(id).scope } + fn expression(&self, id: ScopeId) -> &'a Expression { + self.indexed(id).expression + } + fn select_expression( &self, expression: &Expression, @@ -1399,7 +1425,7 @@ fn to_node( reference_node_name: &str, trim_selects: bool, ) -> Result { - let (context, scope_id) = LineageScopeContext::from_scope(&scope); + let (context, scope_id) = LineageScopeContext::from_scope(&scope, &scope.expression); to_node_inner( column, &context, @@ -1447,7 +1473,7 @@ fn to_node_inner( ))); } let scope = context.scope(scope_id); - let scope_expr = &scope.expression; + let scope_expr = context.expression(scope_id); // Build combined CTE scopes: current scope's cte_scopes + ancestors let mut all_cte_scopes = context.indexed(scope_id).cte_scopes.clone(); @@ -4767,7 +4793,7 @@ mod tests { "WITH base AS (SELECT order_id FROM orders), first_pass AS (SELECT order_id FROM base UNION ALL SELECT order_id FROM base), second_pass AS (SELECT order_id FROM first_pass UNION ALL SELECT order_id FROM first_pass) SELECT order_id FROM second_pass", ); let scope = build_scope(&expression); - let (context, _) = LineageScopeContext::from_scope(&scope); + let (context, _) = LineageScopeContext::from_scope(&scope, &scope.expression); let mut unique_scopes: HashSet<*const Scope> = HashSet::new(); collect_unique_scope_pointers(&scope, &mut unique_scopes); diff --git a/crates/polyglot-sql/src/query_analysis.rs b/crates/polyglot-sql/src/query_analysis.rs index 11f5a1c4..fdb28a99 100644 --- a/crates/polyglot-sql/src/query_analysis.rs +++ b/crates/polyglot-sql/src/query_analysis.rs @@ -858,6 +858,11 @@ struct NullabilityContext<'a> { struct NullabilityScope<'a> { scope: &'a Scope, + /// The scope's complete expression; set-operation branches borrow it from their + /// parent because nested branch scopes store a shallow copy. + expression: &'a Expression, + /// `expression` with query containers unwrapped. + query: &'a Expression, selected: Scope, bindings: HashMap, ctes: HashMap, @@ -896,10 +901,23 @@ impl<'a> NullabilityContext<'a> { resolving: RefCell::new(HashSet::new()), uncertain_columns, }; - context.index_scope(scope, HashMap::new(), None, selected_scopes.as_mut()); + context.index_scope( + scope, + &scope.expression, + HashMap::new(), + None, + selected_scopes.as_mut(), + ); context } + /// The complete expression of an indexed scope, for rooting scoped lineage. + fn scope_expression(&self, scope: &'a Scope) -> &'a Expression { + self.scope_ids + .get(&(scope as *const Scope)) + .map_or(&scope.expression, |id| self.scopes[*id].expression) + } + fn selected_scope(&self, scope: &Scope) -> Option<&Scope> { self.scope_ids .get(&(scope as *const Scope)) @@ -911,6 +929,7 @@ impl<'a> NullabilityContext<'a> { fn index_scope( &mut self, scope: &'a Scope, + expression: &'a Expression, mut ctes: HashMap, recursive_name: Option<&Identifier>, mut selected_scopes: Option<&mut HashMap<*const Scope, Scope>>, @@ -923,7 +942,7 @@ impl<'a> NullabilityContext<'a> { id, ); } - let query = crate::scope::scope_query(&scope.expression); + let query = crate::scope::scope_query(expression); let bindings: HashMap<_, _> = if let Expression::Select(select) = query { select .from @@ -943,6 +962,8 @@ impl<'a> NullabilityContext<'a> { .unwrap_or_else(|| crate::scope::selected_reference_scope(scope)); self.scopes.push(NullabilityScope { scope, + expression, + query, selected, bindings, ctes: HashMap::new(), @@ -950,13 +971,13 @@ impl<'a> NullabilityContext<'a> { branches: Vec::new(), nullable_sources: nullable_source_names(query, self.dialect), }); - let recursive = with_clause(crate::scope::scope_query(&scope.expression)) - .is_some_and(|with| with.recursive); + let recursive = with_clause(query).is_some_and(|with| with.recursive); for child in &scope.cte_scopes { if let Expression::Cte(cte) = &child.expression { let name = &cte.alias; let child_id = self.index_scope( child, + &child.expression, ctes.clone(), recursive.then_some(name), selected_scopes.as_deref_mut(), @@ -969,13 +990,25 @@ impl<'a> NullabilityContext<'a> { } self.scopes[id].ctes = ctes.clone(); for child in &scope.derived_table_scopes { - let child_id = - self.index_scope(child, ctes.clone(), None, selected_scopes.as_deref_mut()); + let child_id = self.index_scope( + child, + &child.expression, + ctes.clone(), + None, + selected_scopes.as_deref_mut(), + ); self.scopes[id].derived.push(child_id); } - for child in &scope.union_scopes { - let child_id = - self.index_scope(child, ctes.clone(), None, selected_scopes.as_deref_mut()); + let branch_queries = crate::scope::set_operation_branch_queries(expression); + for (index, child) in scope.union_scopes.iter().enumerate() { + let branch_query = branch_queries.map_or(&child.expression, |queries| queries[index]); + let child_id = self.index_scope( + child, + branch_query, + ctes.clone(), + None, + selected_scopes.as_deref_mut(), + ); self.scopes[id].branches.push(child_id); } id @@ -1961,8 +1994,9 @@ fn projection_fact<'a>( .as_ref() .and_then(|root| root.for_scope(scope)) .unwrap_or_else(|| { - ScopedLineage::with_schema( + ScopedLineage::with_schema_for_expression( scope, + nullability_context.scope_expression(scope), &[], dialect, nullability_context @@ -2255,8 +2289,9 @@ fn cached_terminal_references_for_expression<'a>( .as_ref() .and_then(|root| root.for_scope(scope)) .unwrap_or_else(|| { - ScopedLineage::with_schema( + ScopedLineage::with_schema_for_expression( scope, + context.scope_expression(scope), &[], dialect, context.mapping_schema.map(|schema| schema as &dyn Schema), @@ -2698,7 +2733,7 @@ impl NullabilityContext<'_> { fn output_inner(&self, scope_id: usize, ordinal: usize, depth: usize) -> ProjectionNullability { use ProjectionNullability::*; let frame = &self.scopes[scope_id]; - let query = crate::scope::scope_query(&frame.scope.expression); + let query = crate::scope::scope_query(frame.query); if let Expression::Select(select) = query { return select .expressions diff --git a/crates/polyglot-sql/src/scope.rs b/crates/polyglot-sql/src/scope.rs index ae20aea6..8b37474f 100644 --- a/crates/polyglot-sql/src/scope.rs +++ b/crates/polyglot-sql/src/scope.rs @@ -239,7 +239,14 @@ pub struct ColumnRef { /// Scopes can be nested (subqueries, CTEs, derived tables) and form a tree. #[derive(Debug, Clone)] pub struct Scope { - /// The expression at the root of this scope + /// The expression at the root of this scope. + /// + /// A set-operation branch that is itself an unmodified set operation (the + /// inner levels of `a UNION b UNION c`) stores its operator and clauses with + /// NULL placeholders for set-operation operands; those operands are the + /// expressions of its own `union_scopes`. Copying them into every level would + /// make long chains quadratic. Borrow a branch's complete query from its + /// parent's with [`set_operation_branch_queries`]. pub expression: Expression, /// Type of this scope relative to its parent @@ -436,7 +443,15 @@ impl Scope { /// This is intended for result schema style output columns (e.g. UNION /// outputs), unlike [`Self::columns`], which returns raw referenced columns. pub fn output_columns(&self) -> Vec { - crate::ast_transforms::get_output_column_names(&self.expression) + // Set-operation output names come from the leftmost branch, which is the + // first branch scope; a nested branch stores only a shallow copy. + match (&self.expression, self.union_scopes.first()) { + ( + Expression::Union(_) | Expression::Intersect(_) | Expression::Except(_), + Some(left), + ) => left.output_columns(), + _ => crate::ast_transforms::get_output_column_names(&self.expression), + } } /// Get all source names in this scope @@ -751,6 +766,79 @@ pub(crate) fn build_scope_with_ctes( root } +/// Borrow the complete queries of a set-operation scope's two branch scopes, +/// in `union_scopes` order, from the scope's complete query. +pub fn set_operation_branch_queries(query: &Expression) -> Option<[&Expression; 2]> { + match scope_query(query) { + Expression::Union(operation) => Some([&operation.left, &operation.right]), + Expression::Intersect(operation) => Some([&operation.left, &operation.right]), + Expression::Except(operation) => Some([&operation.left, &operation.right]), + _ => None, + } +} + +fn is_unmodified_set_operation(expression: &Expression) -> bool { + macro_rules! unmodified { + ($operation:expr) => { + $operation.order_by.is_none() + && $operation.limit.is_none() + && $operation.offset.is_none() + }; + } + match expression { + Expression::Union(operation) => unmodified!(operation), + Expression::Intersect(operation) => unmodified!(operation), + Expression::Except(operation) => unmodified!(operation), + _ => false, + } +} + +/// The expression stored on a set-operation branch scope; see [`Scope::expression`]. +fn set_operation_branch_expression(operand: &Expression) -> Expression { + fn operand_copy(operand: &Expression) -> Expression { + if matches!( + operand, + Expression::Union(_) | Expression::Intersect(_) | Expression::Except(_) + ) { + Expression::Null(crate::expressions::Null) + } else { + operand.clone() + } + } + macro_rules! shallow { + ($variant:ident, $operation:expr) => { + Expression::$variant(Box::new(crate::expressions::$variant { + left: operand_copy(&$operation.left), + right: operand_copy(&$operation.right), + all: $operation.all, + distinct: $operation.distinct, + with: $operation.with.clone(), + order_by: $operation.order_by.clone(), + limit: $operation.limit.clone(), + offset: $operation.offset.clone(), + distribute_by: $operation.distribute_by.clone(), + sort_by: $operation.sort_by.clone(), + cluster_by: $operation.cluster_by.clone(), + by_name: $operation.by_name, + side: $operation.side.clone(), + kind: $operation.kind.clone(), + corresponding: $operation.corresponding, + strict: $operation.strict, + on_columns: $operation.on_columns.clone(), + })) + }; + } + if !is_unmodified_set_operation(operand) { + return operand.clone(); + } + match operand { + Expression::Union(operation) => shallow!(Union, operation), + Expression::Intersect(operation) => shallow!(Intersect, operation), + Expression::Except(operation) => shallow!(Except, operation), + _ => unreachable!("only set operations are shallow"), + } +} + fn build_scope_impl(expression: &Expression, current_scope: &mut Scope) { match expression { Expression::Prepare(prepare) => { @@ -787,11 +875,16 @@ fn build_scope_impl(expression: &Expression, current_scope: &mut Scope) { process_ctes(with, current_scope); } - let mut left_scope = current_scope.branch(union.left.clone(), ScopeType::SetOperation); + let mut left_scope = current_scope.branch( + set_operation_branch_expression(&union.left), + ScopeType::SetOperation, + ); build_scope_impl(&union.left, &mut left_scope); - let mut right_scope = - current_scope.branch(union.right.clone(), ScopeType::SetOperation); + let mut right_scope = current_scope.branch( + set_operation_branch_expression(&union.right), + ScopeType::SetOperation, + ); build_scope_impl(&union.right, &mut right_scope); current_scope.union_scopes.push(left_scope); @@ -802,12 +895,16 @@ fn build_scope_impl(expression: &Expression, current_scope: &mut Scope) { process_ctes(with, current_scope); } - let mut left_scope = - current_scope.branch(intersect.left.clone(), ScopeType::SetOperation); + let mut left_scope = current_scope.branch( + set_operation_branch_expression(&intersect.left), + ScopeType::SetOperation, + ); build_scope_impl(&intersect.left, &mut left_scope); - let mut right_scope = - current_scope.branch(intersect.right.clone(), ScopeType::SetOperation); + let mut right_scope = current_scope.branch( + set_operation_branch_expression(&intersect.right), + ScopeType::SetOperation, + ); build_scope_impl(&intersect.right, &mut right_scope); current_scope.union_scopes.push(left_scope); @@ -818,11 +915,16 @@ fn build_scope_impl(expression: &Expression, current_scope: &mut Scope) { process_ctes(with, current_scope); } - let mut left_scope = current_scope.branch(except.left.clone(), ScopeType::SetOperation); + let mut left_scope = current_scope.branch( + set_operation_branch_expression(&except.left), + ScopeType::SetOperation, + ); build_scope_impl(&except.left, &mut left_scope); - let mut right_scope = - current_scope.branch(except.right.clone(), ScopeType::SetOperation); + let mut right_scope = current_scope.branch( + set_operation_branch_expression(&except.right), + ScopeType::SetOperation, + ); build_scope_impl(&except.right, &mut right_scope); current_scope.union_scopes.push(left_scope); @@ -1823,4 +1925,71 @@ mod tests { "traverse_scope should return scopes for CTAS" ); } + + fn set_operation_chain(branches: usize, operator: &str) -> String { + (0..branches) + .map(|i| format!("SELECT {i} AS id, 'open' AS status FROM orders")) + .collect::>() + .join(&format!(" {operator} ")) + } + + fn stored_expression_nodes(scope: &Scope) -> usize { + scope.expression.dfs().count() + + scope + .union_scopes + .iter() + .chain(&scope.cte_scopes) + .chain(&scope.derived_table_scopes) + .chain(&scope.subquery_scopes) + .map(stored_expression_nodes) + .sum::() + } + + fn set_operation_depth(scope: &Scope) -> usize { + scope + .union_scopes + .first() + .map_or(0, |left| 1 + set_operation_depth(left)) + } + + #[test] + fn long_set_operation_chain_scopes_store_linear_expression_size() { + for operator in ["UNION ALL", "INTERSECT", "EXCEPT"] { + let short = parse_and_build_scope(&set_operation_chain(100, operator)); + let long = parse_and_build_scope(&set_operation_chain(400, operator)); + assert_eq!(set_operation_depth(&short), 99, "{operator}"); + assert_eq!(set_operation_depth(&long), 399, "{operator}"); + let ratio = + stored_expression_nodes(&long) as f64 / stored_expression_nodes(&short) as f64; + assert!( + ratio < 4.5, + "{operator}: 4x branches stored {ratio:.1}x expression nodes" + ); + } + } + + #[test] + fn nested_set_operation_branches_resolve_complete_queries_from_parent() { + let sql = "SELECT id FROM orders UNION ALL BY NAME SELECT id, status FROM returns EXCEPT SELECT id, status FROM orders INTERSECT SELECT id, status FROM returns"; + let ast = Parser::parse_sql(sql).expect("Failed to parse SQL"); + let root = build_scope(&ast[0]); + + fn check(scope: &Scope, query: &Expression) { + assert_eq!( + scope.output_columns(), + crate::ast_transforms::get_output_column_names(query) + ); + let Some(branches) = set_operation_branch_queries(query) else { + assert!(scope.union_scopes.is_empty()); + assert_eq!(&scope.expression, query); + return; + }; + assert_eq!(scope.union_scopes.len(), 2); + for (child, branch) in scope.union_scopes.iter().zip(branches) { + check(child, branch); + } + } + assert_eq!(root.expression, ast[0]); + check(&root, &ast[0]); + } } diff --git a/crates/polyglot-sql/src/validation/tests.rs b/crates/polyglot-sql/src/validation/tests.rs index 172761ff..211c533b 100644 --- a/crates/polyglot-sql/src/validation/tests.rs +++ b/crates/polyglot-sql/src/validation/tests.rs @@ -4529,3 +4529,112 @@ fn set_operation_after_subquery_operand_checks_outer_branch_arity() { ); } } + +fn set_operation_chain_sql( + branches: usize, + operator: &str, + mismatched_branch: Option, +) -> String { + (0..branches) + .map(|i| match mismatched_branch { + Some(branch) if branch == i => format!("SELECT {i} AS id"), + _ => format!("SELECT {i} AS id, CAST({i} AS DOUBLE) AS amount, 'open' AS status"), + }) + .collect::>() + .join(&format!(" {operator} ")) +} + +fn set_operation_chain_options() -> SchemaValidationOptions { + SchemaValidationOptions { + semantic: true, + check_types: true, + check_references: true, + ..Default::default() + } +} + +fn validated_set_operation_chain( + branches: usize, + operator: &str, + mismatched_branch: Option, + dialect: DialectType, +) -> ValidationResult { + validate_with_schema( + &set_operation_chain_sql(branches, operator, mismatched_branch), + dialect, + &semantic_type_schema(), + &set_operation_chain_options(), + ) +} + +#[test] +fn long_set_operation_chain_validation_scales_linearly() { + // Time validation of a parsed chain; parsing is outside this guard. + let fastest = |branches: usize| { + let sql = set_operation_chain_sql(branches, "UNION ALL", None); + let parsed = crate::parse_one(&sql, DialectType::DuckDB).unwrap(); + (0..5) + .map(|_| { + let statement = parsed.clone(); + let start = std::time::Instant::now(); + let result = validate_parsed_with_schema( + vec![statement], + DialectType::DuckDB, + &semantic_type_schema(), + &set_operation_chain_options(), + ); + let elapsed = start.elapsed(); + assert!(result.valid, "{:?}", result.errors); + elapsed + }) + .min() + .unwrap() + }; + fastest(100); + let short = fastest(100); + let long = fastest(400); + let ratio = long.as_secs_f64() / short.as_secs_f64(); + assert!( + ratio < 8.0, + "4x UNION ALL branches took {ratio:.1}x as long ({short:?} -> {long:?})" + ); +} + +#[test] +fn long_set_operation_chain_mismatch_diagnostics_match_short_chain() { + for dialect in [ + DialectType::DuckDB, + DialectType::PostgreSQL, + DialectType::Snowflake, + DialectType::BigQuery, + DialectType::TSQL, + DialectType::Databricks, + ] { + for operator in ["UNION ALL", "EXCEPT", "INTERSECT"] { + for (long_branch, short_branch, counts) in [ + (0, 0, "left 1, right 3"), + (30, 1, "left 3, right 1"), + (59, 2, "left 3, right 1"), + ] { + let long = validated_set_operation_chain(60, operator, Some(long_branch), dialect); + let short = validated_set_operation_chain(3, operator, Some(short_branch), dialect); + assert_eq!( + long.errors.len(), + 1, + "{dialect:?} {operator}: {:?}", + long.errors + ); + assert!( + long.errors[0].message.ends_with(counts), + "{dialect:?} {operator}: {:?}", + long.errors + ); + assert_eq!( + format!("{:?}", long.errors), + format!("{:?}", short.errors), + "{dialect:?} {operator} mismatch at branch {long_branch}" + ); + } + } + } +}