Skip to content

fix: keep set operations after subquery operands at the outer query; build set-operation scopes in linear time #32

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

Merged
kvlonge merged 2 commits into from
Oct 2, 2026
Merged

fix: keep set operations after subquery operands at the outer query; build set-operation scopes in linear time #32

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
7f9c029
fix: keep set operations after subquery operands at the outer query
kvlonge Oct 2, 2026
9ba2a98
perf: build set-operation scopes in linear time
kvlonge Oct 2, 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
4 changes: 2 additions & 2 deletions crates/polyglot-sql/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down
52 changes: 39 additions & 13 deletions crates/polyglot-sql/src/lineage.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<ScopeId>,
derived_table_scopes: Vec<ScopeId>,
cte_scopes: Vec<ScopeId>,
Expand Down Expand Up @@ -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();
Expand All @@ -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));
}
}
}
Expand Down Expand Up @@ -1196,7 +1211,7 @@ impl<'a> ScopedLineage<'a> {

pub(crate) fn output_names(&self) -> Vec<String> {
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,
)
}
Expand Down Expand Up @@ -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(),
Expand All @@ -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;
Expand All @@ -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(),
Expand All @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -1399,7 +1425,7 @@ fn to_node(
reference_node_name: &str,
trim_selects: bool,
) -> Result<LineageNode> {
let (context, scope_id) = LineageScopeContext::from_scope(&scope);
let (context, scope_id) = LineageScopeContext::from_scope(&scope, &scope.expression);
to_node_inner(
column,
&context,
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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);

Expand Down
39 changes: 30 additions & 9 deletions crates/polyglot-sql/src/parser.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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!(
Expand Down Expand Up @@ -31954,7 +31956,17 @@ impl Parser {
}
}

fn parse_parenthesized_primary(&mut self, lparen_comments: Vec<String>) -> Result<Expression> {
/// 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<String>,
set_operation_may_follow: bool,
) -> Result<Expression> {
let raw_start = self.previous().span.start;

// Empty parens () — could be empty tuple or zero-param lambda () -> body
Expand Down Expand Up @@ -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(_)
Expand Down Expand Up @@ -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())?;
Expand Down Expand Up @@ -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
Expand Down
59 changes: 47 additions & 12 deletions crates/polyglot-sql/src/query_analysis.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<String, &'a Expression>,
ctes: HashMap<String, usize>,
Expand Down Expand Up @@ -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))
Expand All @@ -911,6 +929,7 @@ impl<'a> NullabilityContext<'a> {
fn index_scope(
&mut self,
scope: &'a Scope,
expression: &'a Expression,
mut ctes: HashMap<String, usize>,
recursive_name: Option<&Identifier>,
mut selected_scopes: Option<&mut HashMap<*const Scope, Scope>>,
Expand All @@ -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
Expand All @@ -943,20 +962,22 @@ 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(),
derived: Vec::new(),
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(),
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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),
Expand Down Expand Up @@ -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
Expand Down
Loading
Loading