From f43a762d3f777a7c72292d37ea26233e8b2298d6 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Sat, 3 Oct 2026 03:28:42 +0100 Subject: [PATCH 1/2] perf: build comment-wrapped set-operation scopes in linear time A line or block comment before UNION, INTERSECT or EXCEPT makes the parser wrap the preceding operand in a comment annotation. The shallow branch copy from the linear set-operation scope build did not look through that wrapper, so every level of a commented chain copied all earlier branches again and scope building, validation and analysis grew quadratically. Branch scopes now keep the comment wrapper around the shallow copy, and a wrapped set-operation operand gets the same NULL placeholder as a bare one. The scope tree shape and all validation, analysis and lineage results are unchanged. --- crates/polyglot-sql/src/scope.rs | 113 ++++++++++++++++---- crates/polyglot-sql/src/validation/tests.rs | 22 +++- 2 files changed, 110 insertions(+), 25 deletions(-) diff --git a/crates/polyglot-sql/src/scope.rs b/crates/polyglot-sql/src/scope.rs index 8b37474f..01245c33 100644 --- a/crates/polyglot-sql/src/scope.rs +++ b/crates/polyglot-sql/src/scope.rs @@ -242,8 +242,9 @@ pub struct 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 + /// inner levels of `a UNION b UNION c`, including inside the comment wrapper + /// the parser adds before a commented set operator) 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`]. @@ -789,6 +790,19 @@ fn is_unmodified_set_operation(expression: &Expression) -> bool { Expression::Union(operation) => unmodified!(operation), Expression::Intersect(operation) => unmodified!(operation), Expression::Except(operation) => unmodified!(operation), + // The parser wraps a set-operation operand that precedes a commented + // set operator; the wrapper only carries the comments. + Expression::Annotated(annotated) => is_unmodified_set_operation(&annotated.this), + _ => false, + } +} + +/// Whether `operand` is a set operation, possibly inside comment wrappers, +/// and is therefore the expression of its own branch scope. +fn is_set_operation_operand(operand: &Expression) -> bool { + match operand { + Expression::Union(_) | Expression::Intersect(_) | Expression::Except(_) => true, + Expression::Annotated(annotated) => is_set_operation_operand(&annotated.this), _ => false, } } @@ -796,10 +810,7 @@ fn is_unmodified_set_operation(expression: &Expression) -> bool { /// 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(_) - ) { + if is_set_operation_operand(operand) { Expression::Null(crate::expressions::Null) } else { operand.clone() @@ -835,6 +846,12 @@ fn set_operation_branch_expression(operand: &Expression) -> Expression { Expression::Union(operation) => shallow!(Union, operation), Expression::Intersect(operation) => shallow!(Intersect, operation), Expression::Except(operation) => shallow!(Except, operation), + Expression::Annotated(annotated) => { + Expression::Annotated(Box::new(crate::expressions::Annotated { + this: set_operation_branch_expression(&annotated.this), + trailing_comments: annotated.trailing_comments.clone(), + })) + } _ => unreachable!("only set operations are shallow"), } } @@ -1926,11 +1943,13 @@ mod tests { ); } - fn set_operation_chain(branches: usize, operator: &str) -> String { + /// A chain whose operators follow `separator`; a line comment there makes + /// the parser wrap each left operand in a comment annotation. + fn commented_set_operation_chain(branches: usize, operator: &str, separator: &str) -> String { (0..branches) .map(|i| format!("SELECT {i} AS id, 'open' AS status FROM orders")) .collect::>() - .join(&format!(" {operator} ")) + .join(&format!("{separator}{operator} ")) } fn stored_expression_nodes(scope: &Scope) -> usize { @@ -1954,23 +1973,77 @@ mod tests { #[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" - ); + // Commented chains nest twice as deep, so they stay within the parser's + // default depth guard with fewer branches. + for (separator, short_branches) in [(" ", 100), ("\n-- row\n", 50), (" /* row */\n", 50)] { + let long_branches = 4 * short_branches; + for operator in ["UNION ALL", "INTERSECT", "EXCEPT"] { + let short = parse_and_build_scope(&commented_set_operation_chain( + short_branches, + operator, + separator, + )); + let long = parse_and_build_scope(&commented_set_operation_chain( + long_branches, + operator, + separator, + )); + assert_eq!( + set_operation_depth(&short), + short_branches - 1, + "{operator} {separator:?}" + ); + assert_eq!( + set_operation_depth(&long), + long_branches - 1, + "{operator} {separator:?}" + ); + let ratio = + stored_expression_nodes(&long) as f64 / stored_expression_nodes(&short) as f64; + assert!( + ratio < 4.5, + "{operator} {separator:?}: 4x branches stored {ratio:.1}x expression nodes" + ); + } } } + #[test] + fn commented_set_operation_branches_store_comment_wrapped_shallow_copies() { + let sql = commented_set_operation_chain(4, "UNION ALL", "\n-- row\n"); + let ast = Parser::parse_sql(&sql).expect("Failed to parse SQL"); + let root = build_scope(&ast[0]); + assert_eq!(root.expression, ast[0]); + + let Expression::Union(root_union) = &ast[0] else { + panic!("expected a UNION chain"); + }; + let Expression::Annotated(wrapped) = &root_union.left else { + panic!("expected a comment-wrapped left operand"); + }; + let left = &root.union_scopes[0]; + let Expression::Annotated(stored) = &left.expression else { + panic!("expected the branch to keep its comment wrapper"); + }; + assert_eq!(stored.trailing_comments, wrapped.trailing_comments); + let Expression::Union(stored_union) = &stored.this else { + panic!("expected a shallow UNION copy"); + }; + assert!(matches!(stored_union.left, Expression::Null(_))); + assert!(matches!(stored_union.right, Expression::Select(_))); + } + #[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"; + for sql in [ + "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", + "SELECT id FROM orders\n-- first\nUNION ALL BY NAME SELECT id, status FROM returns\n-- second\nEXCEPT SELECT id, status FROM orders /* third */\nINTERSECT SELECT id, status FROM returns", + ] { + check_set_operation_branches(sql); + } + } + + fn check_set_operation_branches(sql: &str) { let ast = Parser::parse_sql(sql).expect("Failed to parse SQL"); let root = build_scope(&ast[0]); diff --git a/crates/polyglot-sql/src/validation/tests.rs b/crates/polyglot-sql/src/validation/tests.rs index 211c533b..a4fc067b 100644 --- a/crates/polyglot-sql/src/validation/tests.rs +++ b/crates/polyglot-sql/src/validation/tests.rs @@ -4569,9 +4569,21 @@ fn validated_set_operation_chain( #[test] fn long_set_operation_chain_validation_scales_linearly() { + assert_set_operation_chain_validation_scales_linearly(" ", 100); +} + +#[test] +fn long_commented_set_operation_chain_validation_scales_linearly() { + // A comment before each operator makes the parser wrap every left operand; + // those chains nest twice as deep, so fewer branches fit the depth guard. + assert_set_operation_chain_validation_scales_linearly("\n-- row\n", 60); +} + +fn assert_set_operation_chain_validation_scales_linearly(separator: &str, short_branches: usize) { // 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 sql = set_operation_chain_sql(branches, "UNION ALL", None) + .replace(" UNION ALL ", &format!("{separator}UNION ALL ")); let parsed = crate::parse_one(&sql, DialectType::DuckDB).unwrap(); (0..5) .map(|_| { @@ -4590,13 +4602,13 @@ fn long_set_operation_chain_validation_scales_linearly() { .min() .unwrap() }; - fastest(100); - let short = fastest(100); - let long = fastest(400); + fastest(short_branches); + let short = fastest(short_branches); + let long = fastest(4 * short_branches); 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:?})" + "4x UNION ALL branches ({separator:?}) took {ratio:.1}x as long ({short:?} -> {long:?})" ); } From 91fa625baaf8023f2065bbe6dfe47814574c4d3f Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Sat, 3 Oct 2026 04:08:45 +0100 Subject: [PATCH 2/2] perf: borrow source column sets in column qualification Schema-aware column qualification copied work proportional to the whole query or relation for every column reference and nesting level: - Each reference copied its source's complete column list, ambiguity checks copied and normalized the column lists of every source, and table lookup copied the unambiguous-column map. The resolver now normalizes each source's columns once and answers these lookups from borrowed sets, and caches correlated outer-scope lookups. - Each SELECT built and copied the complete scope tree of its query, and copied that scope again for nested queries. Qualification now builds only the selected sources of the SELECT and passes enclosing scopes by reference. - Clauses, nested queries, outputs and GROUP BY projections were copied before rewriting; they are now moved, and star expansion returns early when no star is projected. Wide relations and SELECT * over derived tables now qualify in time linear in their width, and nested derived tables no longer rebuild every nested scope at each level. Qualified output and all validation, analysis and lineage results are unchanged. --- .../src/optimizer/qualify_columns.rs | 164 ++++++------ crates/polyglot-sql/src/resolver.rs | 240 +++++++++++------- crates/polyglot-sql/src/scope.rs | 103 ++++++-- .../tests/qualification_scaling.rs | 136 ++++++++++ 4 files changed, 435 insertions(+), 208 deletions(-) create mode 100644 crates/polyglot-sql/tests/qualification_scaling.rs diff --git a/crates/polyglot-sql/src/optimizer/qualify_columns.rs b/crates/polyglot-sql/src/optimizer/qualify_columns.rs index 941cf137..864954cd 100644 --- a/crates/polyglot-sql/src/optimizer/qualify_columns.rs +++ b/crates/polyglot-sql/src/optimizer/qualify_columns.rs @@ -12,9 +12,7 @@ use crate::expressions::{ }; use crate::resolver::{Resolver, ResolverError}; use crate::schema::{normalize_name, Schema}; -use crate::scope::{ - build_scope_with_ctes, selected_reference_scope, traverse_scope, Scope, SourceInfo, SourceKind, -}; +use crate::scope::{selected_source_scope, traverse_scope, Scope, SourceInfo, SourceKind}; use std::cell::RefCell; use std::collections::{HashMap, HashSet}; use thiserror::Error; @@ -173,7 +171,7 @@ fn qualify_nested_queries( schema: &dyn Schema, options: &QualifyColumnsOptions, ctes: &HashMap, - outer: &[Scope], + outer: &[&Scope], ) -> QualifyColumnsResult { transform_in_scope(expression, |node| match node { Expression::Select(_) @@ -189,7 +187,7 @@ fn qualify_query( schema: &dyn Schema, options: &QualifyColumnsOptions, inherited_ctes: &HashMap, - outer: &[Scope], + outer: &[&Scope], ) -> QualifyColumnsResult { #[cfg(feature = "stacker")] { @@ -206,6 +204,11 @@ fn qualify_query( qualify_query_inner(expression, schema, options, inherited_ctes, outer) } +/// Move an expression out of the tree, leaving NULL in its place. +fn take(expression: &mut Expression) -> Expression { + std::mem::replace(expression, Expression::Null(crate::expressions::Null)) +} + fn is_lateral_source(expression: &Expression) -> bool { match expression { Expression::Lateral(_) => true, @@ -221,7 +224,7 @@ fn qualify_query_inner( schema: &dyn Schema, options: &QualifyColumnsOptions, inherited_ctes: &HashMap, - outer: &[Scope], + outer: &[&Scope], ) -> QualifyColumnsResult { let infer_schema = options.infer_schema.unwrap_or(schema.is_empty()); let dialect = options.dialect.or_else(|| schema.dialect()); @@ -279,43 +282,36 @@ fn qualify_query_inner( if let Some(from) = &mut select.from { let sources = std::mem::take(&mut from.expressions); for source in sources { + let lateral_scope = is_lateral_source(&source) + .then(|| selected_source_scope(&Expression::Select(select.clone()), &ctes)); let mut visible_outer = outer.to_vec(); - if is_lateral_source(&source) { - visible_outer.push(selected_reference_scope(&build_scope_with_ctes( - &Expression::Select(select.clone()), - &ctes, - ))); - } + visible_outer.extend(lateral_scope.as_ref()); let qualified = qualify_nested_queries(source, schema, options, &ctes, &visible_outer)?; select.from.as_mut().unwrap().expressions.push(qualified); } } for mut join in joins { - let mut visible_outer = outer.to_vec(); - if is_lateral_source(&join.this) + let lateral_scope = (is_lateral_source(&join.this) || matches!( join.kind, JoinKind::Lateral | JoinKind::LeftLateral | JoinKind::CrossApply | JoinKind::OuterApply - ) - { - visible_outer.push(selected_reference_scope(&build_scope_with_ctes( - &Expression::Select(select.clone()), - &ctes, - ))); - } + )) + .then(|| selected_source_scope(&Expression::Select(select.clone()), &ctes)); + let mut visible_outer = outer.to_vec(); + visible_outer.extend(lateral_scope.as_ref()); join.this = qualify_nested_queries(join.this, schema, options, &ctes, &visible_outer)?; select.joins.push(join); } } if matches!(expression, Expression::Select(_)) { - let scope = selected_reference_scope(&build_scope_with_ctes(&expression, &ctes)); + let scope = selected_source_scope(&expression, &ctes); let mut nested_outer = outer.to_vec(); - nested_outer.push(scope.clone()); + nested_outer.push(&scope); // Temporarily detach relation queries: they were already visited, and // unlike scalar subqueries must not see this SELECT as a parent. let Expression::Select(select) = &mut expression else { @@ -324,9 +320,11 @@ fn qualify_query_inner( let from = select.from.take(); let mut joins = std::mem::take(&mut select.joins); let mut error = None; + // A failure discards the whole expression, so children are moved + // rather than copied. crate::ast_children::for_each_child_mut(&mut expression, |child| { if error.is_none() { - match qualify_nested_queries(child.clone(), schema, options, &ctes, &nested_outer) { + match qualify_nested_queries(take(child), schema, options, &ctes, &nested_outer) { Ok(qualified) => *child = qualified, Err(err) => error = Some(err), } @@ -337,7 +335,7 @@ fn qualify_query_inner( } for join in &mut joins { if let Some(on) = &mut join.on { - *on = qualify_nested_queries(on.clone(), schema, options, &ctes, &nested_outer)?; + *on = qualify_nested_queries(take(on), schema, options, &ctes, &nested_outer)?; } } let Expression::Select(select) = &mut expression else { @@ -370,10 +368,8 @@ fn qualify_query_inner( Expression::Except(query) => (&mut query.left, &mut query.right), _ => unreachable!(), }; - let left_query = std::mem::replace(left, Expression::Null(crate::expressions::Null)); - let right_query = std::mem::replace(right, Expression::Null(crate::expressions::Null)); - *left = qualify_nested_queries(left_query, schema, options, &ctes, outer)?; - *right = qualify_nested_queries(right_query, schema, options, &ctes, outer)?; + *left = qualify_nested_queries(take(left), schema, options, &ctes, outer)?; + *right = qualify_nested_queries(take(right), schema, options, &ctes, outer)?; match &mut expression { Expression::Union(query) => query.with = with, Expression::Intersect(query) => query.with = with, @@ -560,7 +556,8 @@ fn expand_using( column_order: &mut Vec, resolver: &mut Resolver, ) -> bool { - let Ok(source_cols) = resolver.get_source_columns(source_name) else { + let dialect = resolver.dialect; + let Ok(source_cols) = resolver.source_columns(source_name) else { return false; }; let schema_known = !source_cols.is_empty() @@ -568,10 +565,10 @@ fn expand_using( .iter() .any(|column| column == "*" || column.is_empty()); for col_name in source_cols { - let normalized = normalize_column_name(&col_name, resolver.dialect); + let normalized = normalize_column_name(col_name, dialect); if let std::collections::hash_map::Entry::Vacant(entry) = columns.entry(normalized) { entry.insert(source_name.to_string()); - column_order.push(col_name); + column_order.push(col_name.clone()); } } schema_known @@ -690,13 +687,8 @@ fn expand_using( .iter() .filter(|t| { resolver - .get_source_columns(t) - .unwrap_or_default() - .iter() - .any(|column| { - normalize_column_name(column, resolver.dialect) - == normalized_identifier - }) + .source_has_normalized_column(t, &normalized_identifier) + .unwrap_or(false) }) .cloned() .collect(); @@ -828,7 +820,7 @@ fn rewrite_using_columns_in_expression( column_tables: &HashMap>, dialect: Option, ) { - let transformed: crate::Result<_> = transform_in_scope(expr.clone(), |node| match node { + let transformed: crate::Result<_> = transform_in_scope(take(expr), |node| match node { Expression::Column(col) if col.table.is_none() => { let normalized = normalize_column_name(&col.name.name, dialect); if let Some(tables) = column_tables.get(&normalized) { @@ -840,9 +832,7 @@ fn rewrite_using_columns_in_expression( other => Ok(other), }); - if let Ok(next) = transformed { - *expr = next; - } + *expr = transformed.expect("USING rewrite is infallible"); } /// Normalize ambiguous two-part column references into struct/JSON field access. @@ -894,7 +884,7 @@ pub(crate) fn normalize_dotted_columns_in_expression( resolver: &mut Resolver, ) -> QualifyColumnsResult<()> { let resolver = RefCell::new(resolver); - let transformed = transform_in_scope(expression.clone(), |node| { + let transformed = transform_in_scope(take(expression), |node| { let Expression::Column(column) = node else { return Ok(node); }; @@ -928,10 +918,9 @@ pub(crate) fn normalize_dotted_columns_in_expression( field: column.name.clone(), inferred_type: None, }))) - }) - .map_err(|error: crate::Error| QualifyColumnsError::CannotAutoJoin(error.to_string()))?; - - *expression = transformed; + }); + let transformed: Result<_, crate::Error> = transformed; + *expression = transformed.expect("dotted column normalization is infallible"); Ok(()) } @@ -1023,12 +1012,12 @@ fn expand_alias_refs( /// becomes: /// `SELECT a, b FROM t GROUP BY a, b` fn expand_group_by(select: &mut Select, _dialect: Option) -> QualifyColumnsResult<()> { - let projections = select.expressions.clone(); + let projections = &select.expressions; if let Some(group_by) = &mut select.group_by { for group_expr in &mut group_by.expressions { if let Some(index) = positional_reference(group_expr) { - let replacement = select_expression_at_position(&projections, index)?; + let replacement = select_expression_at_position(projections, index)?; *group_expr = replacement; } } @@ -1051,6 +1040,13 @@ fn expand_stars( resolver: &mut Resolver, column_tables: &HashMap>, ) -> QualifyColumnsResult<()> { + if !select.expressions.iter().any(|expr| match expr { + Expression::Star(_) => true, + Expression::Column(col) => is_star_column(col), + _ => false, + }) { + return Ok(()); + } let mut new_selections: Vec = Vec::new(); let mut has_star = false; let mut coalesced_columns: HashSet = HashSet::new(); @@ -1198,25 +1194,25 @@ pub fn qualify_outputs(scope: &Scope) -> QualifyColumnsResult<()> { } fn qualify_outputs_select(select: &mut Select) -> QualifyColumnsResult<()> { - fn qualify_output(expr: &Expression, i: usize) -> Expression { + fn qualify_output(expr: Expression, i: usize) -> Expression { match expr { - Expression::Annotated(annotated) => { - let mut annotated = annotated.clone(); - annotated.this = qualify_output(&annotated.this, i); + Expression::Annotated(mut annotated) => { + annotated.this = qualify_output(annotated.this, i); Expression::Annotated(annotated) } - Expression::Alias(_) => expr.clone(), - Expression::Column(col) => create_alias(expr.clone(), &col.name.name), - Expression::Star(_) => expr.clone(), + Expression::Alias(_) | Expression::Star(_) => expr, + Expression::Column(ref col) => { + let name = col.name.name.clone(); + create_alias(expr, &name) + } _ => { - let alias_name = get_output_name(expr).unwrap_or_else(|| format!("_col_{}", i)); - create_alias(expr.clone(), &alias_name) + let alias_name = get_output_name(&expr).unwrap_or_else(|| format!("_col_{}", i)); + create_alias(expr, &alias_name) } } } - select.expressions = select - .expressions - .iter() + select.expressions = std::mem::take(&mut select.expressions) + .into_iter() .enumerate() .map(|(i, expr)| qualify_output(expr, i)) .collect(); @@ -1232,7 +1228,7 @@ fn qualify_columns_in_expression( let first_error: RefCell> = RefCell::new(None); let resolver_cell: RefCell<&mut Resolver> = RefCell::new(resolver); - let transformed = transform_in_scope(expr.clone(), |node| { + let transformed = transform_in_scope(take(expr), |node| { if first_error.borrow().is_some() { return Ok(node); } @@ -1251,15 +1247,14 @@ fn qualify_columns_in_expression( } _ => Ok(node), } - }) - .map_err(|err: crate::Error| QualifyColumnsError::CannotAutoJoin(err.to_string()))?; - - if let Some(err) = first_error.into_inner() { - return Err(err); + }); + let transformed: Result<_, crate::Error> = transformed; + // On failure the caller discards the whole expression. + *expr = transformed.expect("column qualification transform is infallible"); + match first_error.into_inner() { + Some(err) => Err(err), + None => Ok(()), } - - *expr = transformed; - Ok(()) } fn qualify_single_column( @@ -1274,9 +1269,9 @@ fn qualify_single_column( if let Some(table) = &col.table { let table_name = &table.name; + let normalized_table_name = normalize_name(table_name, resolver.dialect, true, true); let source_name = scope.sources.keys().find(|source| { - normalize_name(source, resolver.dialect, true, true) - == normalize_name(table_name, resolver.dialect, true, true) + normalize_name(source, resolver.dialect, true, true) == normalized_table_name }); let Some(source_name) = source_name else { if resolver.outer_source_exists(table_name) { @@ -1285,16 +1280,17 @@ fn qualify_single_column( return Err(QualifyColumnsError::UnknownTable(table_name.clone())); }; - if let Ok(source_columns) = resolver.get_source_columns(source_name) { + if !allow_partial { let normalized_column_name = normalize_column_name(&col.name.name, resolver.dialect); - if !allow_partial - && !source_columns.is_empty() - && !source_columns.iter().any(|column| { - normalize_column_name(column, resolver.dialect) == normalized_column_name - }) - && !source_columns.contains(&"*".to_string()) + let has_column = + resolver.source_has_normalized_column(source_name, &normalized_column_name); + if let (Ok(false), Ok(source_columns)) = + (has_column, resolver.source_columns(source_name)) { - return Err(QualifyColumnsError::UnknownColumn(col.name.name.clone())); + if !source_columns.is_empty() && !source_columns.iter().any(|column| column == "*") + { + return Err(QualifyColumnsError::UnknownColumn(col.name.name.clone())); + } } } return Ok(()); @@ -1332,7 +1328,7 @@ fn replace_alias_refs_in_expression( alias_to_expression: &HashMap, literal_index: bool, ) { - let transformed: crate::Result<_> = transform_in_scope(expr.clone(), |node| match node { + let transformed: crate::Result<_> = transform_in_scope(take(expr), |node| match node { Expression::Column(col) if col.table.is_none() => { if let Some((alias_expr, index)) = alias_to_expression.get(&col.name.name) { if literal_index && matches!(alias_expr, Expression::Literal(_)) { @@ -1348,9 +1344,7 @@ fn replace_alias_refs_in_expression( other => Ok(other), }); - if let Ok(next) = transformed { - *expr = next; - } + *expr = transformed.expect("alias reference expansion is infallible"); } fn positional_reference(expr: &Expression) -> Option { diff --git a/crates/polyglot-sql/src/resolver.rs b/crates/polyglot-sql/src/resolver.rs index 351af836..7d3e13a3 100644 --- a/crates/polyglot-sql/src/resolver.rs +++ b/crates/polyglot-sql/src/resolver.rs @@ -16,6 +16,7 @@ use crate::generator::Generator; use crate::schema::{normalize_name, Schema}; use crate::scope::{Scope, SourceInfo}; use crate::traversal::ExpressionWalk; +use std::cell::RefCell; use std::collections::{HashMap, HashSet}; use thiserror::Error; @@ -55,7 +56,9 @@ pub struct Resolver<'a> { /// Whether to infer schema from context infer_schema: bool, /// Cached source columns: source_name -> column names - source_columns_cache: HashMap>, + source_columns_cache: HashMap, + /// Sources whose columns resolve, in `scope.sources` order + resolved_sources_cache: Option>, /// Cached unambiguous columns: column_name -> source_name unambiguous_columns_cache: Option>, /// Cached set of all available columns @@ -63,6 +66,45 @@ pub struct Resolver<'a> { /// Actual correlated parents, nearest first. Physical schema entries alone /// do not make a relation visible in a query. outer_scopes: Vec<&'a Scope>, + /// Cached outer-scope column lookups; see `find_column_in_outer_scopes` + outer_columns_cache: RefCell>>, +} + +/// A source's column names and, once lookups repeat, their normalized forms; +/// each is computed once per resolver. +struct SourceColumns { + columns: Vec, + normalized: std::cell::OnceCell>, + /// Whether a lookup has already scanned `columns`. Short-lived resolvers + /// often look up a single column, which a scan answers more cheaply. + scanned: std::cell::Cell, +} + +impl SourceColumns { + fn new(columns: Vec) -> Self { + Self { + columns, + normalized: std::cell::OnceCell::new(), + scanned: std::cell::Cell::new(false), + } + } + + fn contains_normalized(&self, normalized_column: &str, dialect: Option) -> bool { + if self.normalized.get().is_none() && !self.scanned.replace(true) { + return self + .columns + .iter() + .any(|column| normalize_column_name(column, dialect) == normalized_column); + } + self.normalized + .get_or_init(|| { + self.columns + .iter() + .map(|column| normalize_column_name(column, dialect)) + .collect() + }) + .contains(normalized_column) + } } impl<'a> Resolver<'a> { @@ -74,14 +116,16 @@ impl<'a> Resolver<'a> { dialect: schema.dialect(), infer_schema, source_columns_cache: HashMap::new(), + resolved_sources_cache: None, unambiguous_columns_cache: None, all_columns_cache: None, outer_scopes: Vec::new(), + outer_columns_cache: RefCell::new(HashMap::new()), } } - pub(crate) fn with_outer_scopes(mut self, scopes: &'a [Scope]) -> Self { - self.outer_scopes = scopes.iter().rev().collect(); + pub(crate) fn with_outer_scopes(mut self, scopes: &[&'a Scope]) -> Self { + self.outer_scopes = scopes.iter().rev().copied().collect(); self } @@ -95,6 +139,17 @@ impl<'a> Resolver<'a> { } pub(crate) fn find_column_in_outer_scopes(&self, column: &str) -> Option { + if let Some(table) = self.outer_columns_cache.borrow().get(column) { + return table.clone(); + } + let table = self.find_column_in_outer_scopes_uncached(column); + self.outer_columns_cache + .borrow_mut() + .insert(column.to_string(), table.clone()); + table + } + + fn find_column_in_outer_scopes_uncached(&self, column: &str) -> Option { let mut shadowed = HashSet::new(); shadowed.extend( self.scope @@ -123,7 +178,7 @@ impl<'a> Resolver<'a> { /// Returns the table name if it can be found/inferred. pub fn get_table(&mut self, column_name: &str) -> Option { // Try to find table from all sources (unambiguous lookup) - let table_name = self.get_table_name_from_sources(column_name, None); + let table_name = self.get_table_name_from_sources(column_name); // If we found a table, return it if table_name.is_some() { @@ -134,10 +189,11 @@ impl<'a> Resolver<'a> { // assume the column belongs to that source if self.infer_schema { let sources_without_schema: Vec<_> = self - .get_all_source_columns() - .iter() - .filter(|(_, columns)| columns.is_empty() || columns.contains(&"*".to_string())) - .map(|(name, _)| name.clone()) + .resolved_source_columns() + .filter(|(_, source)| { + source.columns.is_empty() || source.columns.iter().any(|column| column == "*") + }) + .map(|(name, _)| name) .collect(); if sources_without_schema.len() == 1 { @@ -182,8 +238,8 @@ impl<'a> Resolver<'a> { pub fn all_columns(&mut self) -> &HashSet { if self.all_columns_cache.is_none() { let mut all = HashSet::new(); - for columns in self.get_all_source_columns().values() { - all.extend(columns.iter().cloned()); + for (_, source) in self.resolved_source_columns() { + all.extend(source.columns.iter().cloned()); } self.all_columns_cache = Some(all); } @@ -196,25 +252,58 @@ impl<'a> Resolver<'a> { /// /// Returns the list of column names available from the given source. pub fn get_source_columns(&mut self, source_name: &str) -> ResolverResult> { - // Check cache first - if let Some(columns) = self.source_columns_cache.get(source_name) { - return Ok(columns.clone()); - } + self.source_columns(source_name).map(<[String]>::to_vec) + } - // Get the source info - let source_info = self - .scope - .sources - .get(source_name) - .ok_or_else(|| ResolverError::UnknownTable(source_name.to_string()))?; + /// Borrow the column names of a source; see [`Self::get_source_columns`]. + pub(crate) fn source_columns(&mut self, source_name: &str) -> ResolverResult<&[String]> { + Ok(&self.cached_source_columns(source_name)?.columns) + } - let columns = self.extract_columns_from_source(source_info)?; + /// Whether a source has a column whose normalized name is `normalized_column`. + pub(crate) fn source_has_normalized_column( + &mut self, + source_name: &str, + normalized_column: &str, + ) -> ResolverResult { + let dialect = self.dialect; + Ok(self + .cached_source_columns(source_name)? + .contains_normalized(normalized_column, dialect)) + } - // Cache the result - self.source_columns_cache - .insert(source_name.to_string(), columns.clone()); + fn cached_source_columns(&mut self, source_name: &str) -> ResolverResult<&SourceColumns> { + if !self.source_columns_cache.contains_key(source_name) { + let source_info = self + .scope + .sources + .get(source_name) + .ok_or_else(|| ResolverError::UnknownTable(source_name.to_string()))?; + let columns = self.extract_columns_from_source(source_info)?; + self.source_columns_cache + .insert(source_name.to_string(), SourceColumns::new(columns)); + } + Ok(&self.source_columns_cache[source_name]) + } - Ok(columns) + /// The scope's sources whose columns resolve, in `scope.sources` order. + fn resolved_source_columns(&mut self) -> impl Iterator { + if self.resolved_sources_cache.is_none() { + let scope = self.scope; + let resolved = scope + .sources + .keys() + .filter(|name| self.cached_source_columns(name).is_ok()) + .cloned() + .collect(); + self.resolved_sources_cache = Some(resolved); + } + let columns = &self.source_columns_cache; + self.resolved_sources_cache + .as_deref() + .expect("cache populated above") + .iter() + .map(move |name| (name, &columns[name.as_str()])) } /// Extract column names from a source expression @@ -571,60 +660,32 @@ impl<'a> Resolver<'a> { } } - /// Get all source columns for all sources in the scope - fn get_all_source_columns(&mut self) -> HashMap> { - let source_names: Vec<_> = self.scope.sources.keys().cloned().collect(); - - let mut result = HashMap::new(); - for source_name in source_names { - if let Ok(columns) = self.get_source_columns(&source_name) { - result.insert(source_name, columns); - } - } - result - } - /// Get the table name for a column from the sources - fn get_table_name_from_sources( - &mut self, - column_name: &str, - source_columns: Option<&HashMap>>, - ) -> Option { + fn get_table_name_from_sources(&mut self, column_name: &str) -> Option { let normalized_column_name = normalize_column_name(column_name, self.dialect); - let unambiguous = match source_columns { - Some(cols) => self.compute_unambiguous_columns(cols), - None => { - if self.unambiguous_columns_cache.is_none() { - let all_source_columns = self.get_all_source_columns(); - self.unambiguous_columns_cache = - Some(self.compute_unambiguous_columns(&all_source_columns)); - } - self.unambiguous_columns_cache - .clone() - .expect("cache populated above") - } - }; - - unambiguous.get(&normalized_column_name).cloned() + if self.unambiguous_columns_cache.is_none() { + let unambiguous = self.compute_unambiguous_columns(); + self.unambiguous_columns_cache = Some(unambiguous); + } + self.unambiguous_columns_cache + .as_ref() + .expect("cache populated above") + .get(&normalized_column_name) + .cloned() } /// Compute unambiguous columns mapping /// - /// A column is unambiguous if it appears in exactly one source. - fn compute_unambiguous_columns( - &self, - source_columns: &HashMap>, - ) -> HashMap { - if source_columns.is_empty() { - return HashMap::new(); - } - + /// A column is unambiguous if it appears in exactly one source. A column + /// listed twice by one source is ambiguous. + fn compute_unambiguous_columns(&mut self) -> HashMap { + let dialect = self.dialect; let mut column_to_sources: HashMap> = HashMap::new(); - for (source_name, columns) in source_columns { - for column in columns { + for (source_name, source) in self.resolved_source_columns() { + for column in &source.columns { column_to_sources - .entry(normalize_column_name(column, self.dialect)) + .entry(normalize_column_name(column, dialect)) .or_default() .push(source_name.clone()); } @@ -641,31 +702,19 @@ impl<'a> Resolver<'a> { /// Check if a column is ambiguous (appears in multiple sources) pub fn is_ambiguous(&mut self, column_name: &str) -> bool { let normalized_column_name = normalize_column_name(column_name, self.dialect); - let all_source_columns = self.get_all_source_columns(); - let sources_with_column: Vec<_> = all_source_columns - .iter() - .filter(|(_, columns)| { - columns.iter().any(|column| { - normalize_column_name(column, self.dialect) == normalized_column_name - }) - }) - .map(|(name, _)| name.clone()) - .collect(); - - sources_with_column.len() > 1 + let dialect = self.dialect; + self.resolved_source_columns() + .filter(|(_, source)| source.contains_normalized(&normalized_column_name, dialect)) + .nth(1) + .is_some() } /// Get all sources that contain a given column pub fn sources_for_column(&mut self, column_name: &str) -> Vec { let normalized_column_name = normalize_column_name(column_name, self.dialect); - let all_source_columns = self.get_all_source_columns(); - all_source_columns - .iter() - .filter(|(_, columns)| { - columns.iter().any(|column| { - normalize_column_name(column, self.dialect) == normalized_column_name - }) - }) + let dialect = self.dialect; + self.resolved_source_columns() + .filter(|(_, source)| source.contains_normalized(&normalized_column_name, dialect)) .map(|(name, _)| name.clone()) .collect() } @@ -683,12 +732,11 @@ impl<'a> Resolver<'a> { let mut matching_sources = Vec::new(); for source_name in available_sources { - if let Ok(columns) = self.get_source_columns(source_name) { - if columns.iter().any(|column| { - normalize_column_name(column, self.dialect) == normalized_column_name - }) { - matching_sources.push(source_name.clone()); - } + if self + .source_has_normalized_column(source_name, &normalized_column_name) + .unwrap_or(false) + { + matching_sources.push(source_name.clone()); } } diff --git a/crates/polyglot-sql/src/scope.rs b/crates/polyglot-sql/src/scope.rs index 01245c33..36a585a6 100644 --- a/crates/polyglot-sql/src/scope.rs +++ b/crates/polyglot-sql/src/scope.rs @@ -113,25 +113,62 @@ pub(crate) fn take_selected_reference_scope( (selected_scope_view(scope, query), with) } -fn selected_scope_view(scope: &Scope, query: Expression) -> Scope { - let aliases: HashSet<_> = walk_in_scope(&query, false) +/// `selected_reference_scope(&build_scope_with_ctes(expression, ctes))` without +/// building child scopes or copying the query. Only `sources` and `cte_sources` +/// are populated; the view's expression is NULL. +pub(crate) fn selected_source_scope( + expression: &Expression, + ctes: &HashMap, +) -> Scope { + let select = match expression { + Expression::Select(select) if select.with.is_none() => select, + _ => return selected_reference_scope(&build_scope_with_ctes(expression, ctes)), + }; + let mut scope = Scope::new(Expression::Null(crate::expressions::Null)); + scope.cte_sources = std::sync::Arc::new(ctes.clone()); + if let Some(from) = &select.from { + for table in &from.expressions { + add_table_to_scope(table, &mut scope, false); + } + } + for join in &select.joins { + add_table_to_scope(&join.this, &mut scope, false); + } + for lateral_view in &select.lateral_views { + add_lateral_view_to_scope(lateral_view, &mut scope); + } + let aliases = selected_table_aliases(expression); + scope + .sources + .retain(|name, source| selected_source(name, source, &aliases)); + scope +} + +fn selected_table_aliases(query: &Expression) -> HashSet { + walk_in_scope(query, false) .filter_map(|node| match node { Expression::Table(table) => { Some(table.alias.as_ref().unwrap_or(&table.name).name.clone()) } _ => None, }) - .collect(); + .collect() +} + +fn selected_source(name: &str, source: &SourceInfo, aliases: &HashSet) -> bool { + // Source keys use the spelling of the actual FROM/JOIN binding. + // Folding here can accidentally retain an unused, quoted CTE. + source.kind != SourceKind::Cte || aliases.contains(name) +} + +fn selected_scope_view(scope: &Scope, query: Expression) -> Scope { + let aliases = selected_table_aliases(&query); let mut selected = Scope::new(query); selected.cte_sources = scope.cte_sources.clone(); selected.sources = scope .sources .iter() - .filter(|(name, source)| { - // Source keys use the spelling of the actual FROM/JOIN binding. - // Folding here can accidentally retain an unused, quoted CTE. - source.kind != SourceKind::Cte || aliases.contains(*name) - }) + .filter(|(name, source)| selected_source(name, source, &aliases)) .map(|(name, source)| (name.clone(), source.clone())) .collect(); selected @@ -870,13 +907,13 @@ fn build_scope_impl(expression: &Expression, current_scope: &mut Scope) { // Process FROM clause if let Some(from) = &select.from { for table in &from.expressions { - add_table_to_scope(table, current_scope); + add_table_to_scope(table, current_scope, true); } } // Process JOINs for join in &select.joins { - add_table_to_scope(&join.this, current_scope); + add_table_to_scope(&join.this, current_scope, true); } // Process table-generating lateral views (Hive/Spark style UDTFs). @@ -999,7 +1036,9 @@ fn cte_body_self_references(cte: &crate::expressions::Cte) -> bool { .is_empty() } -fn add_table_to_scope(expr: &Expression, scope: &mut Scope) { +/// Register a FROM/JOIN relation as a source of `scope`, building the child +/// scopes of derived relations when `build_children` is set. +fn add_table_to_scope(expr: &Expression, scope: &mut Scope, build_children: bool) { match expr { Expression::Table(table) => { let name = table @@ -1041,15 +1080,18 @@ fn add_table_to_scope(expr: &Expression, scope: &mut Scope) { } else { ScopeType::DerivedTable }; - let mut derived_scope = scope.branch(subquery.this.clone(), scope_type); - build_scope_impl(&subquery.this, &mut derived_scope); + let derived_scope = build_children.then(|| { + let mut derived_scope = scope.branch(subquery.this.clone(), scope_type); + build_scope_impl(&subquery.this, &mut derived_scope); + derived_scope + }); if subquery.lateral { scope.add_lateral_source(name, expr.clone(), true); - scope.udtf_scopes.push(derived_scope); + scope.udtf_scopes.extend(derived_scope); } else { scope.add_source(name, expr.clone(), true); - scope.derived_table_scopes.push(derived_scope); + scope.derived_table_scopes.extend(derived_scope); } } Expression::Values(values) => { @@ -1074,17 +1116,20 @@ fn add_table_to_scope(expr: &Expression, scope: &mut Scope) { .iter() .map(|column| column.name.clone()) .collect::>(); - let mut derived_scope = scope.branch_with_options( - alias.this.clone(), - ScopeType::DerivedTable, - None, - None, - Some(outer_columns), - ); - build_scope_impl(&alias.this, &mut derived_scope); + let derived_scope = build_children.then(|| { + let mut derived_scope = scope.branch_with_options( + alias.this.clone(), + ScopeType::DerivedTable, + None, + None, + Some(outer_columns), + ); + build_scope_impl(&alias.this, &mut derived_scope); + derived_scope + }); scope.add_source(alias.alias.name.clone(), expr.clone(), true); - scope.derived_table_scopes.push(derived_scope); + scope.derived_table_scopes.extend(derived_scope); } Expression::Lateral(lateral) => { let name = lateral @@ -1103,7 +1148,9 @@ fn add_table_to_scope(expr: &Expression, scope: &mut Scope) { name, SourceInfo::new(expr.clone(), false, SourceKind::DerivedTable), ); - add_pivot_inner_scope(&pivot.this, scope); + if build_children { + add_pivot_inner_scope(&pivot.this, scope); + } } Expression::Unpivot(unpivot) => { let name = pivot_source_name( @@ -1114,10 +1161,12 @@ fn add_table_to_scope(expr: &Expression, scope: &mut Scope) { name, SourceInfo::new(expr.clone(), false, SourceKind::DerivedTable), ); - add_pivot_inner_scope(&unpivot.this, scope); + if build_children { + add_pivot_inner_scope(&unpivot.this, scope); + } } Expression::Paren(paren) => { - add_table_to_scope(&paren.this, scope); + add_table_to_scope(&paren.this, scope, build_children); } _ => {} } diff --git a/crates/polyglot-sql/tests/qualification_scaling.rs b/crates/polyglot-sql/tests/qualification_scaling.rs new file mode 100644 index 00000000..d8999fc7 --- /dev/null +++ b/crates/polyglot-sql/tests/qualification_scaling.rs @@ -0,0 +1,136 @@ +//! Allocation scaling of schema-aware column qualification. +//! +//! Allocated bytes are deterministic, unlike timings, so these guards compare +//! the work done for an input against the same input at a quarter of the size. +//! This binary holds a single test so no other test allocates concurrently. + +use polyglot_sql::optimizer::{ + qualify_columns, qualify_tables, QualifyColumnsOptions, QualifyTablesOptions, +}; +use polyglot_sql::{parse_one, DataType, DialectType, Expression, MappingSchema, Schema}; +use stats_alloc::{Region, StatsAlloc, INSTRUMENTED_SYSTEM}; +use std::alloc::System; + +#[global_allocator] +static GLOBAL: &StatsAlloc = &INSTRUMENTED_SYSTEM; + +const DIALECT: DialectType = DialectType::Snowflake; + +fn schema(tables: &[&str], width: usize) -> MappingSchema { + let int = DataType::Int { + length: None, + integer_spelling: false, + }; + let mut schema = MappingSchema::new(); + for table in tables { + let columns: Vec<_> = (0..width) + .map(|i| (format!("{table}_c{i}"), int.clone())) + .chain(std::iter::once(("id".to_string(), int.clone()))) + .collect(); + schema.add_table(table, &columns, None).unwrap(); + } + schema +} + +/// Qualify like query analysis and schema-aware lineage do. +fn qualify(expression: Expression, schema: &MappingSchema) -> Expression { + let expression = qualify_tables( + expression, + &QualifyTablesOptions::new() + .with_dialect(DIALECT) + .with_alias_unaliased_tables(false) + .with_alias_unaliased_subqueries(true) + .with_normalize_set_operation_subqueries(false), + ); + qualify_columns( + expression, + schema, + &QualifyColumnsOptions::new() + .with_dialect(DIALECT) + .with_allow_partial(true), + ) + .expect("qualification succeeds") +} + +fn allocated_bytes(sql: &str, schema: &MappingSchema) -> usize { + let parsed = parse_one(sql, DIALECT).expect("query parses"); + let region = Region::new(GLOBAL); + let qualified = qualify(parsed, schema); + let bytes = region.change().bytes_allocated; + drop(qualified); + bytes +} + +/// Bytes allocated for the input built at size 400, relative to size 100. +fn growth(case: impl Fn(usize) -> (String, MappingSchema)) -> f64 { + let (short_sql, short_schema) = case(100); + let (long_sql, long_schema) = case(400); + allocated_bytes(&long_sql, &long_schema) as f64 + / allocated_bytes(&short_sql, &short_schema) as f64 +} + +fn qualified_star_subquery_references(width: usize) -> (String, MappingSchema) { + let references = (0..width) + .map(|i| format!("q.orders_c{i}")) + .collect::>() + .join(", "); + ( + format!("SELECT {references} FROM (SELECT * FROM orders) AS q WHERE q.id > 0"), + schema(&["orders"], width), + ) +} + +fn unqualified_join_references(width: usize) -> (String, MappingSchema) { + let tables = ["orders", "customers", "products", "shipments"]; + let references = tables + .iter() + .flat_map(|table| (0..width).step_by(4).map(move |i| format!("{table}_c{i}"))) + .collect::>() + .join(", "); + ( + format!( + "SELECT {references} FROM orders JOIN customers ON customers.id = orders.id \ + JOIN products ON products.id = orders.id JOIN shipments ON shipments.id = orders.id" + ), + schema(&tables, width), + ) +} + +fn nested_star_subqueries(size: usize) -> (String, MappingSchema) { + let mut query = "SELECT * FROM orders".to_string(); + for level in 0..size / 10 { + query = format!("SELECT * FROM ({query}) AS s{level} WHERE s{level}.id > {level}"); + } + (query, schema(&["orders"], 20)) +} + +#[test] +fn qualification_allocation_scaling() { + // Every column reference used to copy, and for ambiguity checks normalize, + // the complete column lists of its sources. + for (label, case) in [ + ( + "qualified references to a star subquery", + qualified_star_subquery_references as fn(usize) -> (String, MappingSchema), + ), + ( + "unqualified references across joins", + unqualified_join_references, + ), + ] { + let growth = growth(case); + assert!( + growth < 4.5, + "{label}: 4x the width allocated {growth:.1}x the bytes" + ); + } + + // Each nesting level used to rebuild the complete scope tree of every + // relation below it. A level still copies its own derived relation, so the + // work grows quadratically rather than cubically with the depth. + let growth = growth(nested_star_subqueries); + assert!( + growth < 20.0, + "nested star subqueries: 4x the depth allocated {growth:.1}x the bytes" + ); +}