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 8b37474f..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 @@ -242,8 +279,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 +827,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 +847,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 +883,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"), } } @@ -853,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). @@ -982,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 @@ -1024,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) => { @@ -1057,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 @@ -1086,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( @@ -1097,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); } _ => {} } @@ -1926,11 +1992,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 +2022,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:?})" ); } 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" + ); +}