Skip to main content

polyglot_sql/
resolver.rs

1//! Column Resolver Module
2//!
3//! This module provides functionality for resolving column references to their
4//! source tables. It handles:
5//! - Finding which table a column belongs to
6//! - Resolving ambiguous column references
7//! - Handling join context for disambiguation
8//! - Supporting set operations (UNION, INTERSECT, EXCEPT)
9//!
10//! Based on the Python implementation in `sqlglot/optimizer/resolver.py`.
11
12use crate::dialects::DialectType;
13use crate::expressions::{Expression, Identifier, TableRef};
14#[cfg(feature = "generate")]
15use crate::generator::Generator;
16use crate::schema::{normalize_name, Schema};
17use crate::scope::{Scope, SourceInfo};
18use crate::traversal::ExpressionWalk;
19use std::collections::{HashMap, HashSet};
20use thiserror::Error;
21
22/// Errors that can occur during column resolution
23#[derive(Debug, Error, Clone)]
24pub enum ResolverError {
25    #[error("Unknown table: {0}")]
26    UnknownTable(String),
27
28    #[error("Ambiguous column: {column} appears in multiple sources: {sources}")]
29    AmbiguousColumn { column: String, sources: String },
30
31    #[error("Column not found: {0}")]
32    ColumnNotFound(String),
33
34    #[error("Unknown set operation: {0}")]
35    UnknownSetOperation(String),
36
37    #[error("Invalid set operation: {0}")]
38    InvalidSetOperation(String),
39}
40
41/// Result type for resolver operations
42pub type ResolverResult<T> = Result<T, ResolverError>;
43
44/// Helper for resolving columns to their source tables.
45///
46/// This is a struct so we can lazily load some things and easily share
47/// them across functions.
48pub struct Resolver<'a> {
49    /// The scope being analyzed
50    pub scope: &'a Scope,
51    /// The schema for table/column information
52    schema: &'a dyn Schema,
53    /// The dialect being used
54    pub dialect: Option<DialectType>,
55    /// Whether to infer schema from context
56    infer_schema: bool,
57    /// Cached source columns: source_name -> column names
58    source_columns_cache: HashMap<String, Vec<String>>,
59    /// Cached unambiguous columns: column_name -> source_name
60    unambiguous_columns_cache: Option<HashMap<String, String>>,
61    /// Cached set of all available columns
62    all_columns_cache: Option<HashSet<String>>,
63    /// Actual correlated parents, nearest first. Physical schema entries alone
64    /// do not make a relation visible in a query.
65    outer_scopes: Vec<&'a Scope>,
66}
67
68impl<'a> Resolver<'a> {
69    /// Create a new resolver for a scope
70    pub fn new(scope: &'a Scope, schema: &'a dyn Schema, infer_schema: bool) -> Self {
71        Self {
72            scope,
73            schema,
74            dialect: schema.dialect(),
75            infer_schema,
76            source_columns_cache: HashMap::new(),
77            unambiguous_columns_cache: None,
78            all_columns_cache: None,
79            outer_scopes: Vec::new(),
80        }
81    }
82
83    pub(crate) fn with_outer_scopes(mut self, scopes: &'a [Scope]) -> Self {
84        self.outer_scopes = scopes.iter().rev().collect();
85        self
86    }
87
88    pub(crate) fn outer_source_exists(&self, name: &str) -> bool {
89        self.outer_scopes.iter().any(|scope| {
90            scope.sources.keys().any(|source| {
91                normalize_name(source, self.dialect, true, true)
92                    == normalize_name(name, self.dialect, true, true)
93            })
94        })
95    }
96
97    pub(crate) fn find_column_in_outer_scopes(&self, column: &str) -> Option<String> {
98        let mut shadowed = HashSet::new();
99        shadowed.extend(
100            self.scope
101                .sources
102                .keys()
103                .map(|name| normalize_name(name, self.dialect, true, true)),
104        );
105        for outer in &self.outer_scopes {
106            let mut visible = (*outer).clone();
107            visible
108                .sources
109                .retain(|name, _| shadowed.insert(normalize_name(name, self.dialect, true, true)));
110            let mut resolver = Resolver::new(&visible, self.schema, self.infer_schema);
111            if resolver.is_ambiguous(column) {
112                return None;
113            }
114            if let Some(table) = resolver.get_table(column) {
115                return Some(table);
116            }
117        }
118        None
119    }
120
121    /// Get the table for a column name.
122    ///
123    /// Returns the table name if it can be found/inferred.
124    pub fn get_table(&mut self, column_name: &str) -> Option<String> {
125        // Try to find table from all sources (unambiguous lookup)
126        let table_name = self.get_table_name_from_sources(column_name, None);
127
128        // If we found a table, return it
129        if table_name.is_some() {
130            return table_name;
131        }
132
133        // If schema inference is enabled and exactly one source has no schema,
134        // assume the column belongs to that source
135        if self.infer_schema {
136            let sources_without_schema: Vec<_> = self
137                .get_all_source_columns()
138                .iter()
139                .filter(|(_, columns)| columns.is_empty() || columns.contains(&"*".to_string()))
140                .map(|(name, _)| name.clone())
141                .collect();
142
143            if sources_without_schema.len() == 1 {
144                return Some(sources_without_schema[0].clone());
145            }
146        }
147
148        None
149    }
150
151    /// Get the table for a column, returning an Identifier
152    pub fn get_table_identifier(&mut self, column_name: &str) -> Option<Identifier> {
153        self.get_table(column_name).map(Identifier::new)
154    }
155
156    /// Check if a table exists in the schema (not necessarily in the current scope).
157    /// Catalogue membership alone does not imply lexical visibility.
158    pub fn table_exists_in_schema(&self, table_name: &str) -> bool {
159        self.schema.column_names(table_name).is_ok()
160    }
161
162    /// Find the table for a column by searching all schema tables not in the current scope.
163    /// Returns Some(table_name) if the column is found in exactly one non-local
164    /// catalogue table. This is not a correlated-scope lookup: qualification
165    /// must use the actual lexical parents instead.
166    pub fn find_column_in_outer_schema_tables(&self, column_name: &str) -> Option<String> {
167        let tables = self.schema.find_tables_for_column(column_name);
168        // Filter to tables NOT in the current scope
169        let outer_tables: Vec<String> = tables
170            .into_iter()
171            .filter(|t| !self.scope.sources.contains_key(t))
172            .collect();
173        // Only return if unambiguous (exactly one outer table has this column)
174        if outer_tables.len() == 1 {
175            Some(outer_tables.into_iter().next().unwrap())
176        } else {
177            None
178        }
179    }
180
181    /// Get all available columns across all sources in this scope
182    pub fn all_columns(&mut self) -> &HashSet<String> {
183        if self.all_columns_cache.is_none() {
184            let mut all = HashSet::new();
185            for columns in self.get_all_source_columns().values() {
186                all.extend(columns.iter().cloned());
187            }
188            self.all_columns_cache = Some(all);
189        }
190        self.all_columns_cache
191            .as_ref()
192            .expect("cache populated above")
193    }
194
195    /// Get column names for a source.
196    ///
197    /// Returns the list of column names available from the given source.
198    pub fn get_source_columns(&mut self, source_name: &str) -> ResolverResult<Vec<String>> {
199        // Check cache first
200        if let Some(columns) = self.source_columns_cache.get(source_name) {
201            return Ok(columns.clone());
202        }
203
204        // Get the source info
205        let source_info = self
206            .scope
207            .sources
208            .get(source_name)
209            .ok_or_else(|| ResolverError::UnknownTable(source_name.to_string()))?;
210
211        let columns = self.extract_columns_from_source(source_info)?;
212
213        // Cache the result
214        self.source_columns_cache
215            .insert(source_name.to_string(), columns.clone());
216
217        Ok(columns)
218    }
219
220    /// Extract column names from a source expression
221    fn extract_columns_from_source(&self, source_info: &SourceInfo) -> ResolverResult<Vec<String>> {
222        self.get_source_columns_for_expression(&source_info.expression)
223    }
224
225    fn get_source_columns_for_expression(
226        &self,
227        expression: &Expression,
228    ) -> ResolverResult<Vec<String>> {
229        let columns = match expression {
230            Expression::Table(table) => {
231                // For tables, try to get columns from schema.
232                // Build the fully qualified name (catalog.schema.table) to
233                // match how MappingSchema stores hierarchical keys.
234                let table_name = qualified_table_name(table);
235                match self.schema.column_names(&table_name) {
236                    Ok(cols) => cols,
237                    Err(_) => Vec::new(), // Schema might not have this table
238                }
239            }
240            Expression::Subquery(subquery) => {
241                // For subqueries, get named_selects from the inner query
242                apply_alias_columns(
243                    self.get_named_selects(&subquery.this),
244                    &subquery.column_aliases,
245                )
246            }
247            Expression::Values(values) => self.get_values_column_names(values),
248            Expression::Select(select) => {
249                // For derived tables that are SELECT expressions
250                self.get_select_column_names(select)
251            }
252            Expression::Union(union) => {
253                // For UNION, columns come from the set operation
254                self.get_source_columns_from_set_op(&Expression::Union(union.clone()))?
255            }
256            Expression::Intersect(intersect) => {
257                self.get_source_columns_from_set_op(&Expression::Intersect(intersect.clone()))?
258            }
259            Expression::Except(except) => {
260                self.get_source_columns_from_set_op(&Expression::Except(except.clone()))?
261            }
262            Expression::Cte(cte) => {
263                if !cte.columns.is_empty() {
264                    cte.columns.iter().map(|c| c.name.clone()).collect()
265                } else {
266                    self.get_named_selects(&cte.this)
267                }
268            }
269            Expression::Pivot(pivot) => self.get_pivot_output_columns(pivot),
270            Expression::Unpivot(unpivot) => self.get_unpivot_output_columns(unpivot),
271            Expression::Alias(alias) if matches!(&alias.this, Expression::Unnest(_)) => {
272                alias_output_columns(alias)
273            }
274            Expression::Alias(alias) => {
275                let columns = self.get_source_columns_for_expression(&alias.this)?;
276                apply_alias_columns(columns, &alias.column_aliases)
277            }
278            Expression::Unnest(unnest) => unnest_output_columns(unnest),
279            Expression::Lateral(lateral) => lateral_output_columns(lateral),
280            Expression::LateralView(lateral_view) => lateral_view_output_columns(lateral_view),
281            Expression::Paren(paren) => self.get_source_columns_for_expression(&paren.this)?,
282            _ => Vec::new(),
283        };
284
285        Ok(columns)
286    }
287
288    /// Get named selects (column names) from an expression
289    fn get_named_selects(&self, expr: &Expression) -> Vec<String> {
290        if let Ok(identifiers) = crate::set_operation::query_output_identifiers(expr, self.dialect)
291        {
292            return identifiers
293                .into_iter()
294                .map(|identifier| identifier.name)
295                .collect();
296        }
297
298        match expr {
299            Expression::Select(select) => self.get_select_column_names(select),
300            Expression::Union(union) => {
301                // For unions, use the left side's columns
302                self.get_named_selects(&union.left)
303            }
304            Expression::Intersect(intersect) => self.get_named_selects(&intersect.left),
305            Expression::Except(except) => self.get_named_selects(&except.left),
306            Expression::Subquery(subquery) => apply_alias_columns(
307                self.get_named_selects(&subquery.this),
308                &subquery.column_aliases,
309            ),
310            Expression::Values(values) => self.get_values_column_names(values),
311            Expression::Alias(alias) => {
312                let columns = self.get_named_selects(&alias.this);
313                apply_alias_columns(columns, &alias.column_aliases)
314            }
315            Expression::Paren(paren) => self.get_named_selects(&paren.this),
316            _ => Vec::new(),
317        }
318    }
319
320    /// Get column names from a SELECT expression
321    fn get_select_column_names(&self, select: &crate::expressions::Select) -> Vec<String> {
322        select
323            .expressions
324            .iter()
325            .filter_map(|expr| self.get_expression_alias(expr))
326            .collect()
327    }
328
329    fn get_values_column_names(&self, values: &crate::expressions::Values) -> Vec<String> {
330        let width = values
331            .expressions
332            .first()
333            .map_or(0, |row| row.expressions.len());
334        let prefix = match self.dialect {
335            Some(DialectType::Snowflake) => "COLUMN",
336            Some(DialectType::PostgreSQL) => "column",
337            // Default VALUES names are not standardized. Preserve explicit
338            // aliases without inventing names for other dialects.
339            _ => {
340                return values
341                    .column_aliases
342                    .iter()
343                    .map(|alias| alias.name.clone())
344                    .collect()
345            }
346        };
347        apply_alias_columns(
348            (1..=width)
349                .map(|ordinal| format!("{prefix}{ordinal}"))
350                .collect(),
351            &values.column_aliases,
352        )
353    }
354
355    /// Get the alias or name for a select expression
356    fn get_expression_alias(&self, expr: &Expression) -> Option<String> {
357        match expr {
358            Expression::Alias(alias) => Some(alias.alias.name.clone()),
359            Expression::Column(col) => Some(col.name.name.clone()),
360            Expression::Star(_) => Some("*".to_string()),
361            Expression::Identifier(id) => Some(id.name.clone()),
362            _ => None,
363        }
364    }
365
366    fn get_pivot_output_columns(&self, pivot: &crate::expressions::Pivot) -> Vec<String> {
367        if pivot.unpivot {
368            return self.get_pivot_unpivot_output_columns(pivot);
369        }
370
371        let pre_columns = self.get_source_output_columns(&pivot.this);
372        if pre_columns.is_empty() || pre_columns.iter().any(|column| column == "*") {
373            return Vec::new();
374        }
375
376        let excluded = pivot_excluded_source_columns(pivot, self.dialect);
377        let generated = pivot_generated_output_columns(pivot, self.dialect);
378        if excluded.is_empty() || generated.is_empty() {
379            return Vec::new();
380        }
381
382        let mut columns: Vec<String> = pre_columns
383            .into_iter()
384            .filter(|column| !excluded.contains(&normalize_column_name(column, self.dialect)))
385            .collect();
386        columns.extend(generated);
387        apply_alias_columns(columns, &pivot.alias_columns)
388    }
389
390    fn get_pivot_unpivot_output_columns(&self, pivot: &crate::expressions::Pivot) -> Vec<String> {
391        let pre_columns = self.get_source_output_columns(&pivot.this);
392        if pre_columns.is_empty() || pre_columns.iter().any(|column| column == "*") {
393            return Vec::new();
394        }
395
396        let input_columns: HashSet<String> = pivot
397            .expressions
398            .iter()
399            .flat_map(expression_column_names)
400            .map(|column| normalize_column_name(&column, self.dialect))
401            .collect();
402        let mut columns: Vec<String> = pre_columns
403            .into_iter()
404            .filter(|column| !input_columns.contains(&normalize_column_name(column, self.dialect)))
405            .collect();
406
407        if let Some(Expression::UnpivotColumns(unpivot_columns)) = pivot.into.as_deref() {
408            if let Some(name) = expression_name(&unpivot_columns.this) {
409                columns.push(name);
410            }
411            for value_column in &unpivot_columns.expressions {
412                if let Some(name) = expression_name(value_column) {
413                    columns.push(name);
414                }
415            }
416        }
417
418        apply_alias_columns(columns, &pivot.alias_columns)
419    }
420
421    fn get_unpivot_output_columns(&self, unpivot: &crate::expressions::Unpivot) -> Vec<String> {
422        let pre_columns = self.get_source_output_columns(&unpivot.this);
423        if pre_columns.is_empty() || pre_columns.iter().any(|column| column == "*") {
424            return Vec::new();
425        }
426
427        let input_columns: HashSet<String> = unpivot
428            .columns
429            .iter()
430            .flat_map(expression_column_names)
431            .map(|column| normalize_column_name(&column, self.dialect))
432            .collect();
433        let mut columns: Vec<String> = pre_columns
434            .into_iter()
435            .filter(|column| !input_columns.contains(&normalize_column_name(column, self.dialect)))
436            .collect();
437        columns.push(unpivot.name_column.name.clone());
438        columns.push(unpivot.value_column.name.clone());
439        columns.extend(
440            unpivot
441                .extra_value_columns
442                .iter()
443                .map(|column| column.name.clone()),
444        );
445        apply_alias_columns(columns, &unpivot.alias_columns)
446    }
447
448    fn get_source_output_columns(&self, source: &Expression) -> Vec<String> {
449        match source {
450            Expression::Table(table) => {
451                if table.schema.is_none() && table.catalog.is_none() {
452                    if let Some(source) = self.scope.cte_sources.get(&table.name.name) {
453                        return self.extract_columns_from_source(source).unwrap_or_default();
454                    }
455                }
456
457                let table_name = qualified_table_name(table);
458                self.schema.column_names(&table_name).unwrap_or_default()
459            }
460            Expression::Subquery(subquery) => apply_alias_columns(
461                self.get_named_selects(&subquery.this),
462                &subquery.column_aliases,
463            ),
464            Expression::Values(values) => self.get_values_column_names(values),
465            Expression::Select(select) => self.get_select_column_names(select),
466            Expression::Union(_) | Expression::Intersect(_) | Expression::Except(_) => self
467                .get_source_columns_from_set_op(source)
468                .unwrap_or_default(),
469            Expression::Alias(alias) if matches!(&alias.this, Expression::Unnest(_)) => {
470                alias_output_columns(alias)
471            }
472            Expression::Alias(alias) => {
473                let columns = self.get_source_output_columns(&alias.this);
474                apply_alias_columns(columns, &alias.column_aliases)
475            }
476            Expression::Unnest(unnest) => unnest_output_columns(unnest),
477            Expression::Lateral(lateral) => lateral_output_columns(lateral),
478            Expression::LateralView(lateral_view) => lateral_view_output_columns(lateral_view),
479            Expression::Cte(cte) => {
480                if cte.columns.is_empty() {
481                    self.get_named_selects(&cte.this)
482                } else {
483                    cte.columns
484                        .iter()
485                        .map(|column| column.name.clone())
486                        .collect()
487                }
488            }
489            Expression::Paren(paren) => self.get_source_output_columns(&paren.this),
490            _ => Vec::new(),
491        }
492    }
493
494    /// Get columns from a set operation (UNION, INTERSECT, EXCEPT)
495    pub fn get_source_columns_from_set_op(
496        &self,
497        expression: &Expression,
498    ) -> ResolverResult<Vec<String>> {
499        if matches!(
500            expression,
501            Expression::Union(_) | Expression::Intersect(_) | Expression::Except(_)
502        ) {
503            match crate::set_operation::set_operation_layout(expression, self.dialect) {
504                Ok(Some(layout)) => {
505                    return Ok(layout
506                        .outputs
507                        .into_iter()
508                        .map(|output| output.identifier.name)
509                        .collect())
510                }
511                Ok(None) => {}
512                Err(error) if error.is_indeterminate() => {}
513                Err(error) => return Err(ResolverError::InvalidSetOperation(error.to_string())),
514            }
515        }
516
517        match expression {
518            Expression::Select(select) => Ok(self.get_select_column_names(select)),
519            Expression::Subquery(subquery) => {
520                if matches!(
521                    &subquery.this,
522                    Expression::Union(_) | Expression::Intersect(_) | Expression::Except(_)
523                ) {
524                    self.get_source_columns_from_set_op(&subquery.this)
525                } else {
526                    Ok(self.get_named_selects(&subquery.this))
527                }
528            }
529            Expression::Alias(alias) => {
530                let columns = self.get_source_columns_from_set_op(&alias.this)?;
531                Ok(apply_alias_columns(columns, &alias.column_aliases))
532            }
533            Expression::Paren(paren) => self.get_source_columns_from_set_op(&paren.this),
534            Expression::Union(union) => {
535                // Standard UNION: columns come from the left side
536                self.get_source_columns_from_set_op(&union.left)
537            }
538            Expression::Intersect(intersect) => {
539                self.get_source_columns_from_set_op(&intersect.left)
540            }
541            Expression::Except(except) => self.get_source_columns_from_set_op(&except.left),
542            _ => Err(ResolverError::UnknownSetOperation(format!(
543                "{:?}",
544                expression
545            ))),
546        }
547    }
548
549    /// Get all source columns for all sources in the scope
550    fn get_all_source_columns(&mut self) -> HashMap<String, Vec<String>> {
551        let source_names: Vec<_> = self.scope.sources.keys().cloned().collect();
552
553        let mut result = HashMap::new();
554        for source_name in source_names {
555            if let Ok(columns) = self.get_source_columns(&source_name) {
556                result.insert(source_name, columns);
557            }
558        }
559        result
560    }
561
562    /// Get the table name for a column from the sources
563    fn get_table_name_from_sources(
564        &mut self,
565        column_name: &str,
566        source_columns: Option<&HashMap<String, Vec<String>>>,
567    ) -> Option<String> {
568        let normalized_column_name = normalize_column_name(column_name, self.dialect);
569        let unambiguous = match source_columns {
570            Some(cols) => self.compute_unambiguous_columns(cols),
571            None => {
572                if self.unambiguous_columns_cache.is_none() {
573                    let all_source_columns = self.get_all_source_columns();
574                    self.unambiguous_columns_cache =
575                        Some(self.compute_unambiguous_columns(&all_source_columns));
576                }
577                self.unambiguous_columns_cache
578                    .clone()
579                    .expect("cache populated above")
580            }
581        };
582
583        unambiguous.get(&normalized_column_name).cloned()
584    }
585
586    /// Compute unambiguous columns mapping
587    ///
588    /// A column is unambiguous if it appears in exactly one source.
589    fn compute_unambiguous_columns(
590        &self,
591        source_columns: &HashMap<String, Vec<String>>,
592    ) -> HashMap<String, String> {
593        if source_columns.is_empty() {
594            return HashMap::new();
595        }
596
597        let mut column_to_sources: HashMap<String, Vec<String>> = HashMap::new();
598
599        for (source_name, columns) in source_columns {
600            for column in columns {
601                column_to_sources
602                    .entry(normalize_column_name(column, self.dialect))
603                    .or_default()
604                    .push(source_name.clone());
605            }
606        }
607
608        // Keep only columns that appear in exactly one source
609        column_to_sources
610            .into_iter()
611            .filter(|(_, sources)| sources.len() == 1)
612            .map(|(column, sources)| (column, sources.into_iter().next().unwrap()))
613            .collect()
614    }
615
616    /// Check if a column is ambiguous (appears in multiple sources)
617    pub fn is_ambiguous(&mut self, column_name: &str) -> bool {
618        let normalized_column_name = normalize_column_name(column_name, self.dialect);
619        let all_source_columns = self.get_all_source_columns();
620        let sources_with_column: Vec<_> = all_source_columns
621            .iter()
622            .filter(|(_, columns)| {
623                columns.iter().any(|column| {
624                    normalize_column_name(column, self.dialect) == normalized_column_name
625                })
626            })
627            .map(|(name, _)| name.clone())
628            .collect();
629
630        sources_with_column.len() > 1
631    }
632
633    /// Get all sources that contain a given column
634    pub fn sources_for_column(&mut self, column_name: &str) -> Vec<String> {
635        let normalized_column_name = normalize_column_name(column_name, self.dialect);
636        let all_source_columns = self.get_all_source_columns();
637        all_source_columns
638            .iter()
639            .filter(|(_, columns)| {
640                columns.iter().any(|column| {
641                    normalize_column_name(column, self.dialect) == normalized_column_name
642                })
643            })
644            .map(|(name, _)| name.clone())
645            .collect()
646    }
647
648    /// Try to disambiguate a column based on join context
649    ///
650    /// In join conditions, a column can sometimes be disambiguated based on
651    /// which tables have been joined up to that point.
652    pub fn disambiguate_in_join_context(
653        &mut self,
654        column_name: &str,
655        available_sources: &[String],
656    ) -> Option<String> {
657        let normalized_column_name = normalize_column_name(column_name, self.dialect);
658        let mut matching_sources = Vec::new();
659
660        for source_name in available_sources {
661            if let Ok(columns) = self.get_source_columns(source_name) {
662                if columns.iter().any(|column| {
663                    normalize_column_name(column, self.dialect) == normalized_column_name
664                }) {
665                    matching_sources.push(source_name.clone());
666                }
667            }
668        }
669
670        if matching_sources.len() == 1 {
671            Some(matching_sources.remove(0))
672        } else {
673            None
674        }
675    }
676}
677
678fn normalize_column_name(name: &str, dialect: Option<DialectType>) -> String {
679    normalize_name(name, dialect, false, true)
680}
681
682fn apply_alias_columns(mut columns: Vec<String>, alias_columns: &[Identifier]) -> Vec<String> {
683    for (idx, alias) in alias_columns.iter().enumerate() {
684        if let Some(column) = columns.get_mut(idx) {
685            *column = alias.name.clone();
686        }
687    }
688    columns
689}
690
691fn unnest_output_columns(unnest: &crate::expressions::UnnestFunc) -> Vec<String> {
692    unnest
693        .alias
694        .iter()
695        .map(|alias| alias.name.clone())
696        .chain(unnest.offset_alias.iter().map(|alias| alias.name.clone()))
697        .collect()
698}
699
700fn alias_output_columns(alias: &crate::expressions::Alias) -> Vec<String> {
701    if alias.column_aliases.is_empty() {
702        vec![alias.alias.name.clone()]
703    } else {
704        alias
705            .column_aliases
706            .iter()
707            .map(|column| column.name.clone())
708            .collect()
709    }
710}
711
712fn lateral_output_columns(lateral: &crate::expressions::Lateral) -> Vec<String> {
713    if lateral.column_aliases.is_empty() {
714        default_virtual_output_columns(&lateral.this)
715    } else {
716        lateral.column_aliases.clone()
717    }
718}
719
720fn lateral_view_output_columns(lateral_view: &crate::expressions::LateralView) -> Vec<String> {
721    lateral_view
722        .column_aliases
723        .iter()
724        .map(|column| column.name.clone())
725        .collect()
726}
727
728fn default_virtual_output_columns(expression: &Expression) -> Vec<String> {
729    match expression {
730        Expression::Unnest(unnest) => unnest_output_columns(unnest),
731        Expression::Alias(alias) if matches!(&alias.this, Expression::Unnest(_)) => {
732            alias_output_columns(alias)
733        }
734        Expression::Function(function) if function.name.eq_ignore_ascii_case("FLATTEN") => {
735            ["seq", "key", "path", "index", "value", "this"]
736                .into_iter()
737                .map(String::from)
738                .collect()
739        }
740        _ => Vec::new(),
741    }
742}
743
744fn pivot_excluded_source_columns(
745    pivot: &crate::expressions::Pivot,
746    dialect: Option<DialectType>,
747) -> HashSet<String> {
748    pivot
749        .fields
750        .iter()
751        .chain(pivot.expressions.iter())
752        .chain(pivot.using.iter())
753        .flat_map(expression_column_names)
754        .map(|column| normalize_column_name(&column, dialect))
755        .collect()
756}
757
758fn pivot_generated_output_columns(
759    pivot: &crate::expressions::Pivot,
760    _dialect: Option<DialectType>,
761) -> Vec<String> {
762    let fields = pivot_field_output_names(pivot);
763    let aggregations = if pivot.using.is_empty() {
764        &pivot.expressions
765    } else {
766        &pivot.using
767    };
768
769    if fields.is_empty() || aggregations.is_empty() {
770        return Vec::new();
771    }
772
773    let needs_suffix = aggregations.len() > 1;
774    let mut outputs = Vec::new();
775    for field in fields {
776        for aggregation in aggregations {
777            if let Some(suffix) = pivot_aggregation_output_suffix(aggregation, needs_suffix) {
778                outputs.push(format!("{field}_{suffix}"));
779            } else {
780                outputs.push(field.clone());
781            }
782        }
783    }
784    outputs
785}
786
787fn pivot_field_output_names(pivot: &crate::expressions::Pivot) -> Vec<String> {
788    pivot
789        .fields
790        .iter()
791        .filter_map(|field| match field {
792            Expression::In(in_expr) => Some(
793                in_expr
794                    .expressions
795                    .iter()
796                    .filter_map(expression_name)
797                    .collect::<Vec<_>>(),
798            ),
799            _ => None,
800        })
801        .flatten()
802        .collect()
803}
804
805fn pivot_aggregation_output_suffix(expr: &Expression, needs_suffix: bool) -> Option<String> {
806    match expr {
807        Expression::Alias(alias) => Some(alias.alias.name.clone()),
808        _ if needs_suffix => pivot_generated_aggregation_suffix(expr),
809        _ => None,
810    }
811}
812
813#[cfg(feature = "generate")]
814fn pivot_generated_aggregation_suffix(expr: &Expression) -> Option<String> {
815    Generator::sql(expr).ok().map(|sql| sql.to_lowercase())
816}
817
818#[cfg(not(feature = "generate"))]
819fn pivot_generated_aggregation_suffix(expr: &Expression) -> Option<String> {
820    expression_name(expr).or_else(|| Some(expr.variant_name().to_string()))
821}
822
823fn expression_name(expr: &Expression) -> Option<String> {
824    match expr {
825        Expression::PivotAlias(alias) => expression_name(&alias.alias),
826        Expression::Alias(alias) => Some(alias.alias.name.clone()),
827        Expression::Identifier(identifier) => Some(identifier.name.clone()),
828        Expression::Column(column) => Some(column.name.name.clone()),
829        Expression::Literal(literal) => Some(literal.value_str().to_string()),
830        Expression::Var(var) => Some(var.this.clone()),
831        Expression::Tuple(tuple) => tuple.expressions.first().and_then(expression_name),
832        _ => None,
833    }
834}
835
836fn expression_column_names(expr: &Expression) -> Vec<String> {
837    expr.find_all(|node| matches!(node, Expression::Column(_)))
838        .into_iter()
839        .filter_map(|node| match node {
840            Expression::Column(column) => Some(column.name.name.clone()),
841            _ => None,
842        })
843        .collect()
844}
845
846/// Resolve a column to its source table.
847///
848/// This is a convenience function that creates a Resolver and calls get_table.
849pub fn resolve_column(
850    scope: &Scope,
851    schema: &dyn Schema,
852    column_name: &str,
853    infer_schema: bool,
854) -> Option<String> {
855    let mut resolver = Resolver::new(scope, schema, infer_schema);
856    resolver.get_table(column_name)
857}
858
859/// Check if a column is ambiguous in the given scope.
860pub fn is_column_ambiguous(scope: &Scope, schema: &dyn Schema, column_name: &str) -> bool {
861    let mut resolver = Resolver::new(scope, schema, true);
862    resolver.is_ambiguous(column_name)
863}
864
865/// Build the fully qualified table name (catalog.schema.table) from a TableRef.
866fn qualified_table_name(table: &TableRef) -> String {
867    let mut parts = Vec::new();
868    if let Some(catalog) = &table.catalog {
869        parts.push(catalog.name.clone());
870    }
871    if let Some(schema) = &table.schema {
872        parts.push(schema.name.clone());
873    }
874    parts.push(table.name.name.clone());
875    parts.join(".")
876}
877
878#[cfg(test)]
879mod tests {
880    use super::*;
881    use crate::dialects::Dialect;
882    use crate::expressions::DataType;
883    use crate::parser::Parser;
884    use crate::schema::MappingSchema;
885    use crate::scope::build_scope;
886
887    fn create_test_schema() -> MappingSchema {
888        let mut schema = MappingSchema::new();
889        // Add tables with columns
890        schema
891            .add_table(
892                "users",
893                &[
894                    (
895                        "id".to_string(),
896                        DataType::Int {
897                            length: None,
898                            integer_spelling: false,
899                        },
900                    ),
901                    ("name".to_string(), DataType::Text),
902                    ("email".to_string(), DataType::Text),
903                ],
904                None,
905            )
906            .unwrap();
907        schema
908            .add_table(
909                "orders",
910                &[
911                    (
912                        "id".to_string(),
913                        DataType::Int {
914                            length: None,
915                            integer_spelling: false,
916                        },
917                    ),
918                    (
919                        "user_id".to_string(),
920                        DataType::Int {
921                            length: None,
922                            integer_spelling: false,
923                        },
924                    ),
925                    (
926                        "amount".to_string(),
927                        DataType::Double {
928                            precision: None,
929                            scale: None,
930                        },
931                    ),
932                ],
933                None,
934            )
935            .unwrap();
936        schema
937    }
938
939    #[test]
940    fn test_resolver_basic() {
941        let ast = Parser::parse_sql("SELECT id, name FROM users").expect("Failed to parse");
942        let scope = build_scope(&ast[0]);
943        let schema = create_test_schema();
944        let mut resolver = Resolver::new(&scope, &schema, true);
945
946        // 'name' should resolve to 'users' since it's the only source
947        let table = resolver.get_table("name");
948        assert_eq!(table, Some("users".to_string()));
949    }
950
951    #[test]
952    fn test_resolver_ambiguous_column() {
953        let ast =
954            Parser::parse_sql("SELECT id FROM users JOIN orders ON users.id = orders.user_id")
955                .expect("Failed to parse");
956        let scope = build_scope(&ast[0]);
957        let schema = create_test_schema();
958        let mut resolver = Resolver::new(&scope, &schema, true);
959
960        // 'id' appears in both tables, so it's ambiguous
961        assert!(resolver.is_ambiguous("id"));
962
963        // 'name' only appears in users
964        assert!(!resolver.is_ambiguous("name"));
965
966        // 'amount' only appears in orders
967        assert!(!resolver.is_ambiguous("amount"));
968    }
969
970    #[test]
971    fn test_resolver_unambiguous_column() {
972        let ast = Parser::parse_sql(
973            "SELECT name, amount FROM users JOIN orders ON users.id = orders.user_id",
974        )
975        .expect("Failed to parse");
976        let scope = build_scope(&ast[0]);
977        let schema = create_test_schema();
978        let mut resolver = Resolver::new(&scope, &schema, true);
979
980        // 'name' should resolve to 'users'
981        let table = resolver.get_table("name");
982        assert_eq!(table, Some("users".to_string()));
983
984        // 'amount' should resolve to 'orders'
985        let table = resolver.get_table("amount");
986        assert_eq!(table, Some("orders".to_string()));
987    }
988
989    #[test]
990    fn test_resolver_with_alias() {
991        let ast = Parser::parse_sql("SELECT u.id FROM users AS u").expect("Failed to parse");
992        let scope = build_scope(&ast[0]);
993        let schema = create_test_schema();
994        let _resolver = Resolver::new(&scope, &schema, true);
995
996        // Source should be indexed by alias 'u'
997        assert!(scope.sources.contains_key("u"));
998    }
999
1000    #[test]
1001    fn test_sources_for_column() {
1002        let ast = Parser::parse_sql("SELECT * FROM users JOIN orders ON users.id = orders.user_id")
1003            .expect("Failed to parse");
1004        let scope = build_scope(&ast[0]);
1005        let schema = create_test_schema();
1006        let mut resolver = Resolver::new(&scope, &schema, true);
1007
1008        // 'id' should be in both users and orders
1009        let sources = resolver.sources_for_column("id");
1010        assert!(sources.contains(&"users".to_string()));
1011        assert!(sources.contains(&"orders".to_string()));
1012
1013        // 'email' should only be in users
1014        let sources = resolver.sources_for_column("email");
1015        assert_eq!(sources, vec!["users".to_string()]);
1016    }
1017
1018    #[test]
1019    fn test_all_columns() {
1020        let ast = Parser::parse_sql("SELECT * FROM users").expect("Failed to parse");
1021        let scope = build_scope(&ast[0]);
1022        let schema = create_test_schema();
1023        let mut resolver = Resolver::new(&scope, &schema, true);
1024
1025        let all = resolver.all_columns();
1026        assert!(all.contains("id"));
1027        assert!(all.contains("name"));
1028        assert!(all.contains("email"));
1029    }
1030
1031    #[test]
1032    fn test_resolver_cte_projected_alias_column() {
1033        let ast = Parser::parse_sql(
1034            "WITH my_cte AS (SELECT id AS emp_id FROM users) SELECT emp_id FROM my_cte",
1035        )
1036        .expect("Failed to parse");
1037        let scope = build_scope(&ast[0]);
1038        let schema = create_test_schema();
1039        let mut resolver = Resolver::new(&scope, &schema, true);
1040
1041        let table = resolver.get_table("emp_id");
1042        assert_eq!(table, Some("my_cte".to_string()));
1043    }
1044
1045    #[test]
1046    fn test_resolver_uses_name_aligned_set_operation_outputs() {
1047        for dialect_type in [DialectType::DuckDB, DialectType::Snowflake] {
1048            let expr = Dialect::get(dialect_type)
1049                .parse(
1050                    "SELECT combined.right_value FROM (\
1051                     SELECT 1 AS left_value UNION ALL BY NAME \
1052                     SELECT 2 AS right_value\
1053                     ) AS combined",
1054                )
1055                .expect("parse")
1056                .remove(0);
1057            let scope = build_scope(&expr);
1058            let schema = MappingSchema::with_dialect(dialect_type);
1059            let mut resolver = Resolver::new(&scope, &schema, true);
1060
1061            assert_eq!(
1062                resolver.get_table("right_value"),
1063                Some("combined".to_string()),
1064                "right-only output was not exposed for {dialect_type:?}"
1065            );
1066        }
1067
1068        let dialect_type = DialectType::BigQuery;
1069        let expr = Dialect::get(dialect_type)
1070            .parse(
1071                "SELECT combined.c FROM (\
1072                 SELECT 1 AS a, 2 AS b FULL OUTER UNION ALL BY NAME \
1073                 SELECT 3 AS b, 4 AS c\
1074                 ) AS combined",
1075            )
1076            .expect("parse")
1077            .remove(0);
1078        let scope = build_scope(&expr);
1079        let schema = MappingSchema::with_dialect(dialect_type);
1080        let mut resolver = Resolver::new(&scope, &schema, true);
1081        assert_eq!(resolver.get_table("c"), Some("combined".to_string()));
1082    }
1083
1084    #[test]
1085    fn test_resolve_column_helper() {
1086        let ast = Parser::parse_sql("SELECT name FROM users").expect("Failed to parse");
1087        let scope = build_scope(&ast[0]);
1088        let schema = create_test_schema();
1089
1090        let table = resolve_column(&scope, &schema, "name", true);
1091        assert_eq!(table, Some("users".to_string()));
1092    }
1093
1094    #[test]
1095    fn test_resolver_bigquery_mixed_case_column_names() {
1096        let dialect = Dialect::get(DialectType::BigQuery);
1097        let expr = dialect
1098            .parse("SELECT Name AS name FROM teams")
1099            .unwrap()
1100            .into_iter()
1101            .next()
1102            .expect("expected one expression");
1103        let scope = build_scope(&expr);
1104
1105        let mut schema = MappingSchema::with_dialect(DialectType::BigQuery);
1106        schema
1107            .add_table(
1108                "teams",
1109                &[("Name".into(), DataType::String { length: None })],
1110                None,
1111            )
1112            .expect("schema setup");
1113
1114        let mut resolver = Resolver::new(&scope, &schema, true);
1115        let table = resolver.get_table("Name");
1116        assert_eq!(table, Some("teams".to_string()));
1117
1118        let table = resolver.get_table("name");
1119        assert_eq!(table, Some("teams".to_string()));
1120    }
1121}