1use 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#[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
41pub type ResolverResult<T> = Result<T, ResolverError>;
43
44pub struct Resolver<'a> {
49 pub scope: &'a Scope,
51 schema: &'a dyn Schema,
53 pub dialect: Option<DialectType>,
55 infer_schema: bool,
57 source_columns_cache: HashMap<String, Vec<String>>,
59 unambiguous_columns_cache: Option<HashMap<String, String>>,
61 all_columns_cache: Option<HashSet<String>>,
63 outer_scopes: Vec<&'a Scope>,
66}
67
68impl<'a> Resolver<'a> {
69 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 pub fn get_table(&mut self, column_name: &str) -> Option<String> {
125 let table_name = self.get_table_name_from_sources(column_name, None);
127
128 if table_name.is_some() {
130 return table_name;
131 }
132
133 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 pub fn get_table_identifier(&mut self, column_name: &str) -> Option<Identifier> {
153 self.get_table(column_name).map(Identifier::new)
154 }
155
156 pub fn table_exists_in_schema(&self, table_name: &str) -> bool {
159 self.schema.column_names(table_name).is_ok()
160 }
161
162 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 let outer_tables: Vec<String> = tables
170 .into_iter()
171 .filter(|t| !self.scope.sources.contains_key(t))
172 .collect();
173 if outer_tables.len() == 1 {
175 Some(outer_tables.into_iter().next().unwrap())
176 } else {
177 None
178 }
179 }
180
181 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 pub fn get_source_columns(&mut self, source_name: &str) -> ResolverResult<Vec<String>> {
199 if let Some(columns) = self.source_columns_cache.get(source_name) {
201 return Ok(columns.clone());
202 }
203
204 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 self.source_columns_cache
215 .insert(source_name.to_string(), columns.clone());
216
217 Ok(columns)
218 }
219
220 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 let table_name = qualified_table_name(table);
235 match self.schema.column_names(&table_name) {
236 Ok(cols) => cols,
237 Err(_) => Vec::new(), }
239 }
240 Expression::Subquery(subquery) => {
241 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 self.get_select_column_names(select)
251 }
252 Expression::Union(union) => {
253 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 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 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 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 _ => {
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 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 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 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 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 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 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 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 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 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 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
846pub 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
859pub 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
865fn 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 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 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 assert!(resolver.is_ambiguous("id"));
962
963 assert!(!resolver.is_ambiguous("name"));
965
966 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 let table = resolver.get_table("name");
982 assert_eq!(table, Some("users".to_string()));
983
984 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 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 let sources = resolver.sources_for_column("id");
1010 assert!(sources.contains(&"users".to_string()));
1011 assert!(sources.contains(&"orders".to_string()));
1012
1013 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}