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                self.get_named_selects(&subquery.this)
243            }
244            Expression::Select(select) => {
245                // For derived tables that are SELECT expressions
246                self.get_select_column_names(select)
247            }
248            Expression::Union(union) => {
249                // For UNION, columns come from the set operation
250                self.get_source_columns_from_set_op(&Expression::Union(union.clone()))?
251            }
252            Expression::Intersect(intersect) => {
253                self.get_source_columns_from_set_op(&Expression::Intersect(intersect.clone()))?
254            }
255            Expression::Except(except) => {
256                self.get_source_columns_from_set_op(&Expression::Except(except.clone()))?
257            }
258            Expression::Cte(cte) => {
259                if !cte.columns.is_empty() {
260                    cte.columns.iter().map(|c| c.name.clone()).collect()
261                } else {
262                    self.get_named_selects(&cte.this)
263                }
264            }
265            Expression::Pivot(pivot) => self.get_pivot_output_columns(pivot),
266            Expression::Unpivot(unpivot) => self.get_unpivot_output_columns(unpivot),
267            Expression::Alias(alias) if matches!(&alias.this, Expression::Unnest(_)) => {
268                alias_output_columns(alias)
269            }
270            Expression::Alias(alias) => {
271                let columns = self.get_source_columns_for_expression(&alias.this)?;
272                apply_alias_columns(columns, &alias.column_aliases)
273            }
274            Expression::Unnest(unnest) => unnest_output_columns(unnest),
275            Expression::Lateral(lateral) => lateral_output_columns(lateral),
276            Expression::LateralView(lateral_view) => lateral_view_output_columns(lateral_view),
277            Expression::Paren(paren) => self.get_source_columns_for_expression(&paren.this)?,
278            _ => Vec::new(),
279        };
280
281        Ok(columns)
282    }
283
284    /// Get named selects (column names) from an expression
285    fn get_named_selects(&self, expr: &Expression) -> Vec<String> {
286        if let Ok(identifiers) = crate::set_operation::query_output_identifiers(expr, self.dialect)
287        {
288            return identifiers
289                .into_iter()
290                .map(|identifier| identifier.name)
291                .collect();
292        }
293
294        match expr {
295            Expression::Select(select) => self.get_select_column_names(select),
296            Expression::Union(union) => {
297                // For unions, use the left side's columns
298                self.get_named_selects(&union.left)
299            }
300            Expression::Intersect(intersect) => self.get_named_selects(&intersect.left),
301            Expression::Except(except) => self.get_named_selects(&except.left),
302            Expression::Subquery(subquery) => self.get_named_selects(&subquery.this),
303            Expression::Alias(alias) => {
304                let columns = self.get_named_selects(&alias.this);
305                apply_alias_columns(columns, &alias.column_aliases)
306            }
307            Expression::Paren(paren) => self.get_named_selects(&paren.this),
308            _ => Vec::new(),
309        }
310    }
311
312    /// Get column names from a SELECT expression
313    fn get_select_column_names(&self, select: &crate::expressions::Select) -> Vec<String> {
314        select
315            .expressions
316            .iter()
317            .filter_map(|expr| self.get_expression_alias(expr))
318            .collect()
319    }
320
321    /// Get the alias or name for a select expression
322    fn get_expression_alias(&self, expr: &Expression) -> Option<String> {
323        match expr {
324            Expression::Alias(alias) => Some(alias.alias.name.clone()),
325            Expression::Column(col) => Some(col.name.name.clone()),
326            Expression::Star(_) => Some("*".to_string()),
327            Expression::Identifier(id) => Some(id.name.clone()),
328            _ => None,
329        }
330    }
331
332    fn get_pivot_output_columns(&self, pivot: &crate::expressions::Pivot) -> Vec<String> {
333        if pivot.unpivot {
334            return self.get_pivot_unpivot_output_columns(pivot);
335        }
336
337        let pre_columns = self.get_source_output_columns(&pivot.this);
338        if pre_columns.is_empty() || pre_columns.iter().any(|column| column == "*") {
339            return Vec::new();
340        }
341
342        let excluded = pivot_excluded_source_columns(pivot, self.dialect);
343        let generated = pivot_generated_output_columns(pivot, self.dialect);
344        if excluded.is_empty() || generated.is_empty() {
345            return Vec::new();
346        }
347
348        let mut columns: Vec<String> = pre_columns
349            .into_iter()
350            .filter(|column| !excluded.contains(&normalize_column_name(column, self.dialect)))
351            .collect();
352        columns.extend(generated);
353        apply_alias_columns(columns, &pivot.alias_columns)
354    }
355
356    fn get_pivot_unpivot_output_columns(&self, pivot: &crate::expressions::Pivot) -> Vec<String> {
357        let pre_columns = self.get_source_output_columns(&pivot.this);
358        if pre_columns.is_empty() || pre_columns.iter().any(|column| column == "*") {
359            return Vec::new();
360        }
361
362        let input_columns: HashSet<String> = pivot
363            .expressions
364            .iter()
365            .flat_map(expression_column_names)
366            .map(|column| normalize_column_name(&column, self.dialect))
367            .collect();
368        let mut columns: Vec<String> = pre_columns
369            .into_iter()
370            .filter(|column| !input_columns.contains(&normalize_column_name(column, self.dialect)))
371            .collect();
372
373        if let Some(Expression::UnpivotColumns(unpivot_columns)) = pivot.into.as_deref() {
374            if let Some(name) = expression_name(&unpivot_columns.this) {
375                columns.push(name);
376            }
377            for value_column in &unpivot_columns.expressions {
378                if let Some(name) = expression_name(value_column) {
379                    columns.push(name);
380                }
381            }
382        }
383
384        apply_alias_columns(columns, &pivot.alias_columns)
385    }
386
387    fn get_unpivot_output_columns(&self, unpivot: &crate::expressions::Unpivot) -> Vec<String> {
388        let pre_columns = self.get_source_output_columns(&unpivot.this);
389        if pre_columns.is_empty() || pre_columns.iter().any(|column| column == "*") {
390            return Vec::new();
391        }
392
393        let input_columns: HashSet<String> = unpivot
394            .columns
395            .iter()
396            .flat_map(expression_column_names)
397            .map(|column| normalize_column_name(&column, self.dialect))
398            .collect();
399        let mut columns: Vec<String> = pre_columns
400            .into_iter()
401            .filter(|column| !input_columns.contains(&normalize_column_name(column, self.dialect)))
402            .collect();
403        columns.push(unpivot.name_column.name.clone());
404        columns.push(unpivot.value_column.name.clone());
405        columns.extend(
406            unpivot
407                .extra_value_columns
408                .iter()
409                .map(|column| column.name.clone()),
410        );
411        apply_alias_columns(columns, &unpivot.alias_columns)
412    }
413
414    fn get_source_output_columns(&self, source: &Expression) -> Vec<String> {
415        match source {
416            Expression::Table(table) => {
417                if table.schema.is_none() && table.catalog.is_none() {
418                    if let Some(source) = self.scope.cte_sources.get(&table.name.name) {
419                        return self.extract_columns_from_source(source).unwrap_or_default();
420                    }
421                }
422
423                let table_name = qualified_table_name(table);
424                self.schema.column_names(&table_name).unwrap_or_default()
425            }
426            Expression::Subquery(subquery) => self.get_named_selects(&subquery.this),
427            Expression::Select(select) => self.get_select_column_names(select),
428            Expression::Union(_) | Expression::Intersect(_) | Expression::Except(_) => self
429                .get_source_columns_from_set_op(source)
430                .unwrap_or_default(),
431            Expression::Alias(alias) if matches!(&alias.this, Expression::Unnest(_)) => {
432                alias_output_columns(alias)
433            }
434            Expression::Alias(alias) => {
435                let columns = self.get_source_output_columns(&alias.this);
436                apply_alias_columns(columns, &alias.column_aliases)
437            }
438            Expression::Unnest(unnest) => unnest_output_columns(unnest),
439            Expression::Lateral(lateral) => lateral_output_columns(lateral),
440            Expression::LateralView(lateral_view) => lateral_view_output_columns(lateral_view),
441            Expression::Cte(cte) => {
442                if cte.columns.is_empty() {
443                    self.get_named_selects(&cte.this)
444                } else {
445                    cte.columns
446                        .iter()
447                        .map(|column| column.name.clone())
448                        .collect()
449                }
450            }
451            Expression::Paren(paren) => self.get_source_output_columns(&paren.this),
452            _ => Vec::new(),
453        }
454    }
455
456    /// Get columns from a set operation (UNION, INTERSECT, EXCEPT)
457    pub fn get_source_columns_from_set_op(
458        &self,
459        expression: &Expression,
460    ) -> ResolverResult<Vec<String>> {
461        if matches!(
462            expression,
463            Expression::Union(_) | Expression::Intersect(_) | Expression::Except(_)
464        ) {
465            match crate::set_operation::set_operation_layout(expression, self.dialect) {
466                Ok(Some(layout)) => {
467                    return Ok(layout
468                        .outputs
469                        .into_iter()
470                        .map(|output| output.identifier.name)
471                        .collect())
472                }
473                Ok(None) => {}
474                Err(error) if error.is_indeterminate() => {}
475                Err(error) => return Err(ResolverError::InvalidSetOperation(error.to_string())),
476            }
477        }
478
479        match expression {
480            Expression::Select(select) => Ok(self.get_select_column_names(select)),
481            Expression::Subquery(subquery) => {
482                if matches!(
483                    &subquery.this,
484                    Expression::Union(_) | Expression::Intersect(_) | Expression::Except(_)
485                ) {
486                    self.get_source_columns_from_set_op(&subquery.this)
487                } else {
488                    Ok(self.get_named_selects(&subquery.this))
489                }
490            }
491            Expression::Alias(alias) => {
492                let columns = self.get_source_columns_from_set_op(&alias.this)?;
493                Ok(apply_alias_columns(columns, &alias.column_aliases))
494            }
495            Expression::Paren(paren) => self.get_source_columns_from_set_op(&paren.this),
496            Expression::Union(union) => {
497                // Standard UNION: columns come from the left side
498                self.get_source_columns_from_set_op(&union.left)
499            }
500            Expression::Intersect(intersect) => {
501                self.get_source_columns_from_set_op(&intersect.left)
502            }
503            Expression::Except(except) => self.get_source_columns_from_set_op(&except.left),
504            _ => Err(ResolverError::UnknownSetOperation(format!(
505                "{:?}",
506                expression
507            ))),
508        }
509    }
510
511    /// Get all source columns for all sources in the scope
512    fn get_all_source_columns(&mut self) -> HashMap<String, Vec<String>> {
513        let source_names: Vec<_> = self.scope.sources.keys().cloned().collect();
514
515        let mut result = HashMap::new();
516        for source_name in source_names {
517            if let Ok(columns) = self.get_source_columns(&source_name) {
518                result.insert(source_name, columns);
519            }
520        }
521        result
522    }
523
524    /// Get the table name for a column from the sources
525    fn get_table_name_from_sources(
526        &mut self,
527        column_name: &str,
528        source_columns: Option<&HashMap<String, Vec<String>>>,
529    ) -> Option<String> {
530        let normalized_column_name = normalize_column_name(column_name, self.dialect);
531        let unambiguous = match source_columns {
532            Some(cols) => self.compute_unambiguous_columns(cols),
533            None => {
534                if self.unambiguous_columns_cache.is_none() {
535                    let all_source_columns = self.get_all_source_columns();
536                    self.unambiguous_columns_cache =
537                        Some(self.compute_unambiguous_columns(&all_source_columns));
538                }
539                self.unambiguous_columns_cache
540                    .clone()
541                    .expect("cache populated above")
542            }
543        };
544
545        unambiguous.get(&normalized_column_name).cloned()
546    }
547
548    /// Compute unambiguous columns mapping
549    ///
550    /// A column is unambiguous if it appears in exactly one source.
551    fn compute_unambiguous_columns(
552        &self,
553        source_columns: &HashMap<String, Vec<String>>,
554    ) -> HashMap<String, String> {
555        if source_columns.is_empty() {
556            return HashMap::new();
557        }
558
559        let mut column_to_sources: HashMap<String, Vec<String>> = HashMap::new();
560
561        for (source_name, columns) in source_columns {
562            for column in columns {
563                column_to_sources
564                    .entry(normalize_column_name(column, self.dialect))
565                    .or_default()
566                    .push(source_name.clone());
567            }
568        }
569
570        // Keep only columns that appear in exactly one source
571        column_to_sources
572            .into_iter()
573            .filter(|(_, sources)| sources.len() == 1)
574            .map(|(column, sources)| (column, sources.into_iter().next().unwrap()))
575            .collect()
576    }
577
578    /// Check if a column is ambiguous (appears in multiple sources)
579    pub fn is_ambiguous(&mut self, column_name: &str) -> bool {
580        let normalized_column_name = normalize_column_name(column_name, self.dialect);
581        let all_source_columns = self.get_all_source_columns();
582        let sources_with_column: Vec<_> = all_source_columns
583            .iter()
584            .filter(|(_, columns)| {
585                columns.iter().any(|column| {
586                    normalize_column_name(column, self.dialect) == normalized_column_name
587                })
588            })
589            .map(|(name, _)| name.clone())
590            .collect();
591
592        sources_with_column.len() > 1
593    }
594
595    /// Get all sources that contain a given column
596    pub fn sources_for_column(&mut self, column_name: &str) -> Vec<String> {
597        let normalized_column_name = normalize_column_name(column_name, self.dialect);
598        let all_source_columns = self.get_all_source_columns();
599        all_source_columns
600            .iter()
601            .filter(|(_, columns)| {
602                columns.iter().any(|column| {
603                    normalize_column_name(column, self.dialect) == normalized_column_name
604                })
605            })
606            .map(|(name, _)| name.clone())
607            .collect()
608    }
609
610    /// Try to disambiguate a column based on join context
611    ///
612    /// In join conditions, a column can sometimes be disambiguated based on
613    /// which tables have been joined up to that point.
614    pub fn disambiguate_in_join_context(
615        &mut self,
616        column_name: &str,
617        available_sources: &[String],
618    ) -> Option<String> {
619        let normalized_column_name = normalize_column_name(column_name, self.dialect);
620        let mut matching_sources = Vec::new();
621
622        for source_name in available_sources {
623            if let Ok(columns) = self.get_source_columns(source_name) {
624                if columns.iter().any(|column| {
625                    normalize_column_name(column, self.dialect) == normalized_column_name
626                }) {
627                    matching_sources.push(source_name.clone());
628                }
629            }
630        }
631
632        if matching_sources.len() == 1 {
633            Some(matching_sources.remove(0))
634        } else {
635            None
636        }
637    }
638}
639
640fn normalize_column_name(name: &str, dialect: Option<DialectType>) -> String {
641    normalize_name(name, dialect, false, true)
642}
643
644fn apply_alias_columns(mut columns: Vec<String>, alias_columns: &[Identifier]) -> Vec<String> {
645    for (idx, alias) in alias_columns.iter().enumerate() {
646        if let Some(column) = columns.get_mut(idx) {
647            *column = alias.name.clone();
648        }
649    }
650    columns
651}
652
653fn unnest_output_columns(unnest: &crate::expressions::UnnestFunc) -> Vec<String> {
654    unnest
655        .alias
656        .iter()
657        .map(|alias| alias.name.clone())
658        .chain(unnest.offset_alias.iter().map(|alias| alias.name.clone()))
659        .collect()
660}
661
662fn alias_output_columns(alias: &crate::expressions::Alias) -> Vec<String> {
663    if alias.column_aliases.is_empty() {
664        vec![alias.alias.name.clone()]
665    } else {
666        alias
667            .column_aliases
668            .iter()
669            .map(|column| column.name.clone())
670            .collect()
671    }
672}
673
674fn lateral_output_columns(lateral: &crate::expressions::Lateral) -> Vec<String> {
675    if lateral.column_aliases.is_empty() {
676        default_virtual_output_columns(&lateral.this)
677    } else {
678        lateral.column_aliases.clone()
679    }
680}
681
682fn lateral_view_output_columns(lateral_view: &crate::expressions::LateralView) -> Vec<String> {
683    lateral_view
684        .column_aliases
685        .iter()
686        .map(|column| column.name.clone())
687        .collect()
688}
689
690fn default_virtual_output_columns(expression: &Expression) -> Vec<String> {
691    match expression {
692        Expression::Unnest(unnest) => unnest_output_columns(unnest),
693        Expression::Alias(alias) if matches!(&alias.this, Expression::Unnest(_)) => {
694            alias_output_columns(alias)
695        }
696        Expression::Function(function) if function.name.eq_ignore_ascii_case("FLATTEN") => {
697            ["seq", "key", "path", "index", "value", "this"]
698                .into_iter()
699                .map(String::from)
700                .collect()
701        }
702        _ => Vec::new(),
703    }
704}
705
706fn pivot_excluded_source_columns(
707    pivot: &crate::expressions::Pivot,
708    dialect: Option<DialectType>,
709) -> HashSet<String> {
710    pivot
711        .fields
712        .iter()
713        .chain(pivot.expressions.iter())
714        .chain(pivot.using.iter())
715        .flat_map(expression_column_names)
716        .map(|column| normalize_column_name(&column, dialect))
717        .collect()
718}
719
720fn pivot_generated_output_columns(
721    pivot: &crate::expressions::Pivot,
722    _dialect: Option<DialectType>,
723) -> Vec<String> {
724    let fields = pivot_field_output_names(pivot);
725    let aggregations = if pivot.using.is_empty() {
726        &pivot.expressions
727    } else {
728        &pivot.using
729    };
730
731    if fields.is_empty() || aggregations.is_empty() {
732        return Vec::new();
733    }
734
735    let needs_suffix = aggregations.len() > 1;
736    let mut outputs = Vec::new();
737    for field in fields {
738        for aggregation in aggregations {
739            if let Some(suffix) = pivot_aggregation_output_suffix(aggregation, needs_suffix) {
740                outputs.push(format!("{field}_{suffix}"));
741            } else {
742                outputs.push(field.clone());
743            }
744        }
745    }
746    outputs
747}
748
749fn pivot_field_output_names(pivot: &crate::expressions::Pivot) -> Vec<String> {
750    pivot
751        .fields
752        .iter()
753        .filter_map(|field| match field {
754            Expression::In(in_expr) => Some(
755                in_expr
756                    .expressions
757                    .iter()
758                    .filter_map(expression_name)
759                    .collect::<Vec<_>>(),
760            ),
761            _ => None,
762        })
763        .flatten()
764        .collect()
765}
766
767fn pivot_aggregation_output_suffix(expr: &Expression, needs_suffix: bool) -> Option<String> {
768    match expr {
769        Expression::Alias(alias) => Some(alias.alias.name.clone()),
770        _ if needs_suffix => pivot_generated_aggregation_suffix(expr),
771        _ => None,
772    }
773}
774
775#[cfg(feature = "generate")]
776fn pivot_generated_aggregation_suffix(expr: &Expression) -> Option<String> {
777    Generator::sql(expr).ok().map(|sql| sql.to_lowercase())
778}
779
780#[cfg(not(feature = "generate"))]
781fn pivot_generated_aggregation_suffix(expr: &Expression) -> Option<String> {
782    expression_name(expr).or_else(|| Some(expr.variant_name().to_string()))
783}
784
785fn expression_name(expr: &Expression) -> Option<String> {
786    match expr {
787        Expression::PivotAlias(alias) => expression_name(&alias.alias),
788        Expression::Alias(alias) => Some(alias.alias.name.clone()),
789        Expression::Identifier(identifier) => Some(identifier.name.clone()),
790        Expression::Column(column) => Some(column.name.name.clone()),
791        Expression::Literal(literal) => Some(literal.value_str().to_string()),
792        Expression::Var(var) => Some(var.this.clone()),
793        Expression::Tuple(tuple) => tuple.expressions.first().and_then(expression_name),
794        _ => None,
795    }
796}
797
798fn expression_column_names(expr: &Expression) -> Vec<String> {
799    expr.find_all(|node| matches!(node, Expression::Column(_)))
800        .into_iter()
801        .filter_map(|node| match node {
802            Expression::Column(column) => Some(column.name.name.clone()),
803            _ => None,
804        })
805        .collect()
806}
807
808/// Resolve a column to its source table.
809///
810/// This is a convenience function that creates a Resolver and calls get_table.
811pub fn resolve_column(
812    scope: &Scope,
813    schema: &dyn Schema,
814    column_name: &str,
815    infer_schema: bool,
816) -> Option<String> {
817    let mut resolver = Resolver::new(scope, schema, infer_schema);
818    resolver.get_table(column_name)
819}
820
821/// Check if a column is ambiguous in the given scope.
822pub fn is_column_ambiguous(scope: &Scope, schema: &dyn Schema, column_name: &str) -> bool {
823    let mut resolver = Resolver::new(scope, schema, true);
824    resolver.is_ambiguous(column_name)
825}
826
827/// Build the fully qualified table name (catalog.schema.table) from a TableRef.
828fn qualified_table_name(table: &TableRef) -> String {
829    let mut parts = Vec::new();
830    if let Some(catalog) = &table.catalog {
831        parts.push(catalog.name.clone());
832    }
833    if let Some(schema) = &table.schema {
834        parts.push(schema.name.clone());
835    }
836    parts.push(table.name.name.clone());
837    parts.join(".")
838}
839
840#[cfg(test)]
841mod tests {
842    use super::*;
843    use crate::dialects::Dialect;
844    use crate::expressions::DataType;
845    use crate::parser::Parser;
846    use crate::schema::MappingSchema;
847    use crate::scope::build_scope;
848
849    fn create_test_schema() -> MappingSchema {
850        let mut schema = MappingSchema::new();
851        // Add tables with columns
852        schema
853            .add_table(
854                "users",
855                &[
856                    (
857                        "id".to_string(),
858                        DataType::Int {
859                            length: None,
860                            integer_spelling: false,
861                        },
862                    ),
863                    ("name".to_string(), DataType::Text),
864                    ("email".to_string(), DataType::Text),
865                ],
866                None,
867            )
868            .unwrap();
869        schema
870            .add_table(
871                "orders",
872                &[
873                    (
874                        "id".to_string(),
875                        DataType::Int {
876                            length: None,
877                            integer_spelling: false,
878                        },
879                    ),
880                    (
881                        "user_id".to_string(),
882                        DataType::Int {
883                            length: None,
884                            integer_spelling: false,
885                        },
886                    ),
887                    (
888                        "amount".to_string(),
889                        DataType::Double {
890                            precision: None,
891                            scale: None,
892                        },
893                    ),
894                ],
895                None,
896            )
897            .unwrap();
898        schema
899    }
900
901    #[test]
902    fn test_resolver_basic() {
903        let ast = Parser::parse_sql("SELECT id, name FROM users").expect("Failed to parse");
904        let scope = build_scope(&ast[0]);
905        let schema = create_test_schema();
906        let mut resolver = Resolver::new(&scope, &schema, true);
907
908        // 'name' should resolve to 'users' since it's the only source
909        let table = resolver.get_table("name");
910        assert_eq!(table, Some("users".to_string()));
911    }
912
913    #[test]
914    fn test_resolver_ambiguous_column() {
915        let ast =
916            Parser::parse_sql("SELECT id FROM users JOIN orders ON users.id = orders.user_id")
917                .expect("Failed to parse");
918        let scope = build_scope(&ast[0]);
919        let schema = create_test_schema();
920        let mut resolver = Resolver::new(&scope, &schema, true);
921
922        // 'id' appears in both tables, so it's ambiguous
923        assert!(resolver.is_ambiguous("id"));
924
925        // 'name' only appears in users
926        assert!(!resolver.is_ambiguous("name"));
927
928        // 'amount' only appears in orders
929        assert!(!resolver.is_ambiguous("amount"));
930    }
931
932    #[test]
933    fn test_resolver_unambiguous_column() {
934        let ast = Parser::parse_sql(
935            "SELECT name, amount FROM users JOIN orders ON users.id = orders.user_id",
936        )
937        .expect("Failed to parse");
938        let scope = build_scope(&ast[0]);
939        let schema = create_test_schema();
940        let mut resolver = Resolver::new(&scope, &schema, true);
941
942        // 'name' should resolve to 'users'
943        let table = resolver.get_table("name");
944        assert_eq!(table, Some("users".to_string()));
945
946        // 'amount' should resolve to 'orders'
947        let table = resolver.get_table("amount");
948        assert_eq!(table, Some("orders".to_string()));
949    }
950
951    #[test]
952    fn test_resolver_with_alias() {
953        let ast = Parser::parse_sql("SELECT u.id FROM users AS u").expect("Failed to parse");
954        let scope = build_scope(&ast[0]);
955        let schema = create_test_schema();
956        let _resolver = Resolver::new(&scope, &schema, true);
957
958        // Source should be indexed by alias 'u'
959        assert!(scope.sources.contains_key("u"));
960    }
961
962    #[test]
963    fn test_sources_for_column() {
964        let ast = Parser::parse_sql("SELECT * FROM users JOIN orders ON users.id = orders.user_id")
965            .expect("Failed to parse");
966        let scope = build_scope(&ast[0]);
967        let schema = create_test_schema();
968        let mut resolver = Resolver::new(&scope, &schema, true);
969
970        // 'id' should be in both users and orders
971        let sources = resolver.sources_for_column("id");
972        assert!(sources.contains(&"users".to_string()));
973        assert!(sources.contains(&"orders".to_string()));
974
975        // 'email' should only be in users
976        let sources = resolver.sources_for_column("email");
977        assert_eq!(sources, vec!["users".to_string()]);
978    }
979
980    #[test]
981    fn test_all_columns() {
982        let ast = Parser::parse_sql("SELECT * FROM users").expect("Failed to parse");
983        let scope = build_scope(&ast[0]);
984        let schema = create_test_schema();
985        let mut resolver = Resolver::new(&scope, &schema, true);
986
987        let all = resolver.all_columns();
988        assert!(all.contains("id"));
989        assert!(all.contains("name"));
990        assert!(all.contains("email"));
991    }
992
993    #[test]
994    fn test_resolver_cte_projected_alias_column() {
995        let ast = Parser::parse_sql(
996            "WITH my_cte AS (SELECT id AS emp_id FROM users) SELECT emp_id FROM my_cte",
997        )
998        .expect("Failed to parse");
999        let scope = build_scope(&ast[0]);
1000        let schema = create_test_schema();
1001        let mut resolver = Resolver::new(&scope, &schema, true);
1002
1003        let table = resolver.get_table("emp_id");
1004        assert_eq!(table, Some("my_cte".to_string()));
1005    }
1006
1007    #[test]
1008    fn test_resolver_uses_name_aligned_set_operation_outputs() {
1009        for dialect_type in [DialectType::DuckDB, DialectType::Snowflake] {
1010            let expr = Dialect::get(dialect_type)
1011                .parse(
1012                    "SELECT combined.right_value FROM (\
1013                     SELECT 1 AS left_value UNION ALL BY NAME \
1014                     SELECT 2 AS right_value\
1015                     ) AS combined",
1016                )
1017                .expect("parse")
1018                .remove(0);
1019            let scope = build_scope(&expr);
1020            let schema = MappingSchema::with_dialect(dialect_type);
1021            let mut resolver = Resolver::new(&scope, &schema, true);
1022
1023            assert_eq!(
1024                resolver.get_table("right_value"),
1025                Some("combined".to_string()),
1026                "right-only output was not exposed for {dialect_type:?}"
1027            );
1028        }
1029
1030        let dialect_type = DialectType::BigQuery;
1031        let expr = Dialect::get(dialect_type)
1032            .parse(
1033                "SELECT combined.c FROM (\
1034                 SELECT 1 AS a, 2 AS b FULL OUTER UNION ALL BY NAME \
1035                 SELECT 3 AS b, 4 AS c\
1036                 ) AS combined",
1037            )
1038            .expect("parse")
1039            .remove(0);
1040        let scope = build_scope(&expr);
1041        let schema = MappingSchema::with_dialect(dialect_type);
1042        let mut resolver = Resolver::new(&scope, &schema, true);
1043        assert_eq!(resolver.get_table("c"), Some("combined".to_string()));
1044    }
1045
1046    #[test]
1047    fn test_resolve_column_helper() {
1048        let ast = Parser::parse_sql("SELECT name FROM users").expect("Failed to parse");
1049        let scope = build_scope(&ast[0]);
1050        let schema = create_test_schema();
1051
1052        let table = resolve_column(&scope, &schema, "name", true);
1053        assert_eq!(table, Some("users".to_string()));
1054    }
1055
1056    #[test]
1057    fn test_resolver_bigquery_mixed_case_column_names() {
1058        let dialect = Dialect::get(DialectType::BigQuery);
1059        let expr = dialect
1060            .parse("SELECT Name AS name FROM teams")
1061            .unwrap()
1062            .into_iter()
1063            .next()
1064            .expect("expected one expression");
1065        let scope = build_scope(&expr);
1066
1067        let mut schema = MappingSchema::with_dialect(DialectType::BigQuery);
1068        schema
1069            .add_table(
1070                "teams",
1071                &[("Name".into(), DataType::String { length: None })],
1072                None,
1073            )
1074            .expect("schema setup");
1075
1076        let mut resolver = Resolver::new(&scope, &schema, true);
1077        let table = resolver.get_table("Name");
1078        assert_eq!(table, Some("teams".to_string()));
1079
1080        let table = resolver.get_table("name");
1081        assert_eq!(table, Some("teams".to_string()));
1082    }
1083}