From 6a9832db098c1517a5253bf4d0b49828540ecc4f Mon Sep 17 00:00:00 2001 From: osipovartem Date: Fri, 23 Oct 2026 19:47:00 +0300 Subject: [PATCH] Resolve Snowflake WHERE aliases against input columns --- datafusion/sql/src/query.rs | 2 +- datafusion/sql/src/select.rs | 346 +++++++++++++++++++++++- datafusion/sql/tests/sql_integration.rs | 212 +++++++++++++++ 3 files changed, 550 insertions(+), 10 deletions(-) diff --git a/datafusion/sql/src/query.rs b/datafusion/sql/src/query.rs index 583c04ec516f4..5e5b968d90f10 100644 --- a/datafusion/sql/src/query.rs +++ b/datafusion/sql/src/query.rs @@ -122,7 +122,7 @@ impl SqlToRel<'_, S> { ) -> Result { match pipe_operator { PipeOperator::Where { expr } => { - self.plan_selection(Some(expr), plan, planner_context) + self.plan_selection(Some(expr), plan, &[], planner_context) } PipeOperator::OrderBy { exprs } => { let sort_exprs = self.order_by_to_sort_expr( diff --git a/datafusion/sql/src/select.rs b/datafusion/sql/src/select.rs index 10e6ba0925a44..3a14d08d7152c 100644 --- a/datafusion/sql/src/select.rs +++ b/datafusion/sql/src/select.rs @@ -15,11 +15,11 @@ // specific language governing permissions and limitations // under the License. -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::ops::ControlFlow; use std::sync::Arc; -use crate::planner::{ContextProvider, PlannerContext, SqlToRel}; +use crate::planner::{ContextProvider, IdentNormalizer, PlannerContext, SqlToRel}; use crate::query::to_order_by_exprs_with_select; use crate::utils::{ CheckColumnsMustReferenceAggregatePurpose, CheckColumnsSatisfyExprsPurpose, @@ -33,7 +33,9 @@ use arrow::datatypes::DataType; use datafusion_common::config::Dialect as SqlDialect; use datafusion_common::error::DataFusionErrorBuilder; use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; -use datafusion_common::{Column, DFSchema, DFSchemaRef, Result, not_impl_err, plan_err}; +use datafusion_common::{ + Column, DFSchema, DFSchemaRef, Result, not_impl_err, plan_datafusion_err, plan_err, +}; use datafusion_common::{NullHandling, RecursionUnnestOption, UnnestOptions}; use datafusion_expr::ExprSchemable; use datafusion_expr::builder::get_struct_unnested_columns; @@ -54,11 +56,12 @@ use datafusion_expr::{ use indexmap::IndexMap; use sqlparser::ast::{ - Distinct, ExcludeSelectItem, Expr as SQLExpr, GroupByExpr, NamedWindowExpr, - ObjectNamePart, OrderBy, RenameSelectItem, SelectItemQualifiedWildcardKind, - WildcardAdditionalOptions, WindowType, visit_expressions_mut, + Distinct, ExcludeSelectItem, Expr as SQLExpr, GroupByExpr, Ident, NamedWindowExpr, + ObjectNamePart, OneOrManyWithParens, OrderBy, RenameSelectItem, + SelectItemQualifiedWildcardKind, VisitMut, VisitorMut, WildcardAdditionalOptions, + WindowType, visit_expressions, visit_expressions_mut, }; -use sqlparser::ast::{NamedWindowDefinition, Select, SelectItem, TableWithJoins}; +use sqlparser::ast::{NamedWindowDefinition, Query, Select, SelectItem, TableWithJoins}; /// Result of the `aggregate` function, containing the aggregate plan and /// rewritten expressions that reference the aggregate output columns. @@ -89,6 +92,303 @@ struct RewrittenUnnestExprGroups { expr_groups: Vec>, } +const MAX_WHERE_ALIAS_EXPANSIONS: usize = 256; +const MAX_WHERE_ALIAS_NODES: usize = 4096; + +struct WhereAliasResolver<'a> { + normalizer: IdentNormalizer, + input_schema: &'a DFSchema, + projection: &'a [SelectItem], + aliases: HashMap>, + expanded: HashMap, + active: HashSet, + replacements: usize, + replacement_nodes: usize, + query_depth: usize, +} + +impl<'a> WhereAliasResolver<'a> { + fn new( + normalizer: IdentNormalizer, + input_schema: &'a DFSchema, + projection: &'a [SelectItem], + ) -> Self { + let mut aliases: HashMap> = HashMap::new(); + for item in projection { + if let SelectItem::ExprWithAlias { expr, alias } = item { + let name = normalizer.normalize(alias.clone()); + aliases + .entry(name) + .and_modify(|existing| *existing = None) + .or_insert(Some(expr)); + } + } + Self { + normalizer, + input_schema, + projection, + aliases, + expanded: HashMap::new(), + active: HashSet::new(), + replacements: 0, + replacement_nodes: 0, + query_depth: 0, + } + } + + fn resolve(&mut self, expr: &mut SQLExpr) -> Result<()> { + match expr.visit(self) { + ControlFlow::Continue(()) => Ok(()), + ControlFlow::Break(error) => Err(error), + } + } + + fn resolve_alias(&mut self, name: &str) -> Result { + if let Some(expr) = self.expanded.get(name) { + return Ok(expr.clone()); + } + let Some(alias) = self.aliases.get(name).copied().flatten() else { + return plan_err!("Ambiguous SELECT alias in WHERE: {name}"); + }; + if !self.active.insert(name.to_string()) { + return plan_err!("Circular SELECT alias in WHERE: {name}"); + } + let mut expr = alias.clone(); + if contains_lambda(&expr) { + WhereLambdaAlphaRenamer::new( + self.normalizer, + self.input_schema, + self.projection, + &expr, + ) + .rename(&mut expr)?; + } + self.resolve(&mut expr)?; + self.active.remove(name); + self.expanded.insert(name.to_string(), expr.clone()); + Ok(expr) + } + + fn count_replacement_nodes(&mut self, expr: &SQLExpr) -> Result<()> { + let remaining = MAX_WHERE_ALIAS_NODES - self.replacement_nodes; + let mut nodes = 0; + let result = visit_expressions(expr, |_: &SQLExpr| { + nodes += 1; + if nodes > remaining { + ControlFlow::Break(()) + } else { + ControlFlow::Continue(()) + } + }); + if result.is_break() { + return plan_err!( + "WHERE alias expansion exceeds {MAX_WHERE_ALIAS_NODES} expression nodes" + ); + } + self.replacement_nodes += nodes; + Ok(()) + } +} + +fn contains_lambda(expr: &SQLExpr) -> bool { + visit_expressions(expr, |node| { + if matches!(node, SQLExpr::Lambda(_)) { + ControlFlow::Break(()) + } else { + ControlFlow::Continue(()) + } + }) + .is_break() +} + +struct WhereLambdaAlphaRenamer { + normalizer: IdentNormalizer, + relevant_names: HashSet, + used_names: HashSet, + scopes: Vec>, + query_boundaries: Vec, + next_name: usize, +} + +impl WhereLambdaAlphaRenamer { + fn new( + normalizer: IdentNormalizer, + input_schema: &DFSchema, + projection: &[SelectItem], + predicate: &SQLExpr, + ) -> Self { + let mut relevant_names = HashSet::new(); + let mut used_names = HashSet::new(); + for item in projection { + if let SelectItem::ExprWithAlias { expr, alias } = item { + let name = normalizer.normalize(alias.clone()); + relevant_names.insert(name.clone()); + used_names.insert(name); + let _ = visit_expressions(expr, |node| { + Self::collect_names(normalizer, node, &mut relevant_names); + Self::collect_names(normalizer, node, &mut used_names); + ControlFlow::<()>::Continue(()) + }); + } + } + let _ = visit_expressions(predicate, |node| { + Self::collect_names(normalizer, node, &mut used_names); + ControlFlow::<()>::Continue(()) + }); + for (_, field) in input_schema.iter() { + used_names.insert(field.name().to_string()); + } + Self { + normalizer, + relevant_names, + used_names, + scopes: Vec::new(), + query_boundaries: Vec::new(), + next_name: 0, + } + } + + fn collect_names( + normalizer: IdentNormalizer, + node: &SQLExpr, + names: &mut HashSet, + ) { + if let SQLExpr::Identifier(ident) = node { + names.insert(normalizer.normalize(ident.clone())); + } + if let SQLExpr::Lambda(lambda) = node { + names.extend( + lambda + .params + .iter() + .map(|param| normalizer.normalize(param.name.clone())), + ); + } + } + + fn fresh_name(&mut self) -> Ident { + loop { + let name = format!("__DF_WHERE_ALIAS_LAMBDA_{}", self.next_name); + self.next_name += 1; + if self.used_names.insert(name.clone()) { + return Ident::with_quote('"', name); + } + } + } + + fn rename(&mut self, expr: &mut SQLExpr) -> Result<()> { + match expr.visit(self) { + ControlFlow::Continue(()) => Ok(()), + ControlFlow::Break(error) => Err(error), + } + } +} + +impl VisitorMut for WhereLambdaAlphaRenamer { + type Break = datafusion_common::DataFusionError; + + fn pre_visit_expr(&mut self, expr: &mut SQLExpr) -> ControlFlow { + if let SQLExpr::Lambda(lambda) = expr { + let mut scope = HashMap::new(); + let mut seen = HashSet::new(); + let params = match &mut lambda.params { + OneOrManyWithParens::One(param) => std::slice::from_mut(param), + OneOrManyWithParens::Many(params) => params.as_mut_slice(), + }; + for param in params { + let name = self.normalizer.normalize(param.name.clone()); + if !seen.insert(name.clone()) { + return ControlFlow::Break(plan_datafusion_err!( + "Duplicate lambda parameter: {name}" + )); + } + if self.relevant_names.contains(&name) { + let replacement = self.fresh_name(); + scope.insert(name, replacement.clone()); + param.name = replacement; + } + } + self.scopes.push(scope); + } + ControlFlow::Continue(()) + } + + fn pre_visit_query(&mut self, _query: &mut Query) -> ControlFlow { + self.query_boundaries.push(self.scopes.len()); + ControlFlow::Continue(()) + } + + fn post_visit_query(&mut self, _query: &mut Query) -> ControlFlow { + self.query_boundaries.pop(); + ControlFlow::Continue(()) + } + + fn post_visit_expr(&mut self, expr: &mut SQLExpr) -> ControlFlow { + if matches!(expr, SQLExpr::Lambda(_)) { + self.scopes.pop(); + return ControlFlow::Continue(()); + } + let SQLExpr::Identifier(ident) = expr else { + return ControlFlow::Continue(()); + }; + let name = self.normalizer.normalize(ident.clone()); + let start = self.query_boundaries.last().copied().unwrap_or(0); + if let Some(replacement) = self.scopes[start..] + .iter() + .rev() + .find_map(|scope| scope.get(&name)) + { + *ident = replacement.clone(); + } + ControlFlow::Continue(()) + } +} + +impl VisitorMut for WhereAliasResolver<'_> { + type Break = datafusion_common::DataFusionError; + + fn pre_visit_query(&mut self, _query: &mut Query) -> ControlFlow { + self.query_depth += 1; + ControlFlow::Continue(()) + } + + fn post_visit_query(&mut self, _query: &mut Query) -> ControlFlow { + self.query_depth -= 1; + ControlFlow::Continue(()) + } + + fn post_visit_expr(&mut self, expr: &mut SQLExpr) -> ControlFlow { + if self.query_depth != 0 { + return ControlFlow::Continue(()); + } + let SQLExpr::Identifier(identifier) = expr else { + return ControlFlow::Continue(()); + }; + let name = self.normalizer.normalize(identifier.clone()); + if self.input_schema.has_column_with_unqualified_name(&name) + || !self.aliases.contains_key(&name) + { + return ControlFlow::Continue(()); + } + if self.replacements >= MAX_WHERE_ALIAS_EXPANSIONS { + return ControlFlow::Break(plan_datafusion_err!( + "WHERE alias expansion exceeds {MAX_WHERE_ALIAS_EXPANSIONS} references" + )); + } + self.replacements += 1; + match self.resolve_alias(&name) { + Ok(replacement) => { + if let Err(error) = self.count_replacement_nodes(&replacement) { + return ControlFlow::Break(error); + } + *expr = replacement; + ControlFlow::Continue(()) + } + Err(error) => ControlFlow::Break(error), + } + } +} + fn contains_nested_window_function(expr: &Expr) -> Result { let mut found = false; expr.apply_children(|child| { @@ -190,7 +490,12 @@ impl SqlToRel<'_, S> { let empty_from = matches!(plan, LogicalPlan::EmptyRelation(_)); // Process `where` clause - let base_plan = self.plan_selection(select.selection, plan, planner_context)?; + let base_plan = self.plan_selection( + select.selection, + plan, + &select.projection, + planner_context, + )?; // Handle named windows before processing the projection expression check_conflicting_windows(&select.named_window)?; @@ -1047,13 +1352,36 @@ impl SqlToRel<'_, S> { &self, selection: Option, plan: LogicalPlan, + projection: &[SelectItem], planner_context: &mut PlannerContext, ) -> Result { match selection { - Some(predicate_expr) => { + Some(mut predicate_expr) => { let fallback_schemas = plan.fallback_normalize_schemas(); self.warn_on_null_equality_predicate(&predicate_expr); + if self.context_provider.options().sql_parser.dialect + == SqlDialect::Snowflake + && projection + .iter() + .any(|item| matches!(item, SelectItem::ExprWithAlias { .. })) + { + if contains_lambda(&predicate_expr) { + WhereLambdaAlphaRenamer::new( + self.ident_normalizer, + plan.schema(), + projection, + &predicate_expr, + ) + .rename(&mut predicate_expr)?; + } + WhereAliasResolver::new( + self.ident_normalizer, + plan.schema(), + projection, + ) + .resolve(&mut predicate_expr)?; + } let filter_expr = self.sql_to_expr(predicate_expr, plan.schema(), planner_context)?; diff --git a/datafusion/sql/tests/sql_integration.rs b/datafusion/sql/tests/sql_integration.rs index bc22434ea651e..167cfafb7ae53 100644 --- a/datafusion/sql/tests/sql_integration.rs +++ b/datafusion/sql/tests/sql_integration.rs @@ -5187,6 +5187,218 @@ fn logical_plan(sql: &str) -> Result { logical_plan_with_options(sql, ParserOptions::default()) } +#[test] +fn snowflake_where_aliases_respect_input_columns() { + let mut config = datafusion_common::config::ConfigOptions::default(); + config.sql_parser.dialect = datafusion_common::config::Dialect::Snowflake; + + let plan = logical_plan_with_config( + "SELECT 1 AS age FROM person WHERE age = 42", + config.clone(), + ) + .unwrap(); + assert_contains!( + plan.display_indent().to_string(), + "Filter: person.age = Int64(42)" + ); + + let plan = logical_plan_with_config( + "SELECT age + 1 AS next_age FROM person WHERE next_age > 42", + config.clone(), + ) + .unwrap(); + assert_contains!( + plan.display_indent().to_string(), + "Filter: person.age + Int64(1) > Int64(42)" + ); + + let plan = logical_plan_with_config( + "SELECT age + 1 AS age FROM person WHERE age > 42", + config.clone(), + ) + .unwrap(); + assert_contains!( + plan.display_indent().to_string(), + "Filter: person.age > Int64(42)" + ); + + let plan = logical_plan_with_config( + "SELECT age + 1 AS \"NextAge\" FROM person WHERE \"NextAge\" + 1 > 42", + config.clone(), + ) + .unwrap(); + assert_contains!( + plan.display_indent().to_string(), + "Filter: person.age + Int64(1) + Int64(1) > Int64(42)" + ); + + let plan = logical_plan_with_config( + "WITH src AS (SELECT age FROM person) SELECT 1 AS age FROM src WHERE age > 42", + config.clone(), + ) + .unwrap(); + assert_contains!( + plan.display_indent().to_string(), + "Filter: src.age > Int64(42)" + ); + + let plan = logical_plan_with_config( + "SELECT 1 AS x FROM person WHERE EXISTS (SELECT age AS x FROM person p WHERE x > 42)", + config.clone(), + ) + .unwrap(); + assert_contains!( + plan.display_indent().to_string(), + "Filter: p.age > Int64(42)" + ); + + let plan = logical_plan_with_config("SELECT 123 AS a WHERE a > 100", config).unwrap(); + assert_contains!( + plan.display_indent().to_string(), + "Filter: Int64(123) > Int64(100)" + ); +} + +#[test] +fn snowflake_where_aliases_take_precedence_over_outer_columns() { + let mut config = datafusion_common::config::ConfigOptions::default(); + config.sql_parser.dialect = datafusion_common::config::Dialect::Snowflake; + + let plan = logical_plan_with_config( + "SELECT (SELECT 1 AS age WHERE age = 1) FROM person", + config, + ) + .unwrap(); + assert_contains!(plan.display_indent().to_string(), "Int64(1) = Int64(1)"); +} + +#[test] +fn snowflake_where_aliases_keep_ambiguous_and_invalid_references() { + let mut config = datafusion_common::config::ConfigOptions::default(); + config.sql_parser.dialect = datafusion_common::config::Dialect::Snowflake; + + let err = logical_plan_with_config( + "SELECT 1 AS age FROM person p JOIN person q ON p.id = q.id WHERE age = 42", + config.clone(), + ) + .unwrap_err(); + assert_contains!(err.to_string(), "Ambiguous"); + + let err = logical_plan_with_config("SELECT a AS a WHERE a = 1", config.clone()) + .unwrap_err(); + assert_contains!(err.to_string(), "Circular SELECT alias in WHERE"); + + let err = logical_plan_with_config( + "SELECT SUM(age) AS total FROM person WHERE total > 42", + config, + ) + .unwrap_err(); + assert_contains!( + err.to_string(), + "Aggregate functions are not allowed in the WHERE clause" + ); + + let err = logical_plan("SELECT 123 AS a WHERE a > 100").unwrap_err(); + assert_contains!(err.to_string(), "No field named a"); +} + +#[test] +fn snowflake_where_aliases_do_not_capture_lambda_parameters() { + let mut state = MockSessionState::default() + .with_higher_order_function(Arc::new(HigherOrderUDF::new_from_impl( + MockArrayReduce::new(), + ))) + .with_scalar_function(make_array_udf()) + .with_expr_planner(Arc::new(CustomExprPlanner {})); + state.config_options.sql_parser.dialect = + datafusion_common::config::Dialect::Snowflake; + + let plan = logical_plan_from_state( + "SELECT 999 AS acc WHERE array_reduce([1.0, 2.0], 0, (acc, v) -> acc + v, v -> -v) > 0", + &DatabricksDialect, + ParserOptions::default(), + state, + ) + .unwrap(); + assert_contains!( + plan.display_indent().to_string(), + "__DF_WHERE_ALIAS_LAMBDA_0 + v" + ); + + let mut state = MockSessionState::default() + .with_higher_order_function(Arc::new(HigherOrderUDF::new_from_impl( + MockArrayReduce::new(), + ))) + .with_scalar_function(make_array_udf()) + .with_expr_planner(Arc::new(CustomExprPlanner {})); + state.config_options.sql_parser.dialect = + datafusion_common::config::Dialect::Snowflake; + let plan = logical_plan_from_state( + "SELECT 5 AS k WHERE array_reduce([1.0, 2.0], 0, (acc, v) -> acc + k, v -> -v) > 0", + &DatabricksDialect, + ParserOptions::default(), + state, + ) + .unwrap(); + assert_contains!(plan.display_indent().to_string(), "acc + Int64(5)"); +} + +#[test] +fn snowflake_where_aliases_in_lambda_preserve_column_sources() { + let mut state = MockSessionState::default() + .with_higher_order_function(Arc::new(HigherOrderUDF::new_from_impl( + MockArrayReduce::new(), + ))) + .with_scalar_function(make_array_udf()) + .with_expr_planner(Arc::new(CustomExprPlanner {})); + state.config_options.sql_parser.dialect = + datafusion_common::config::Dialect::Snowflake; + let plan = logical_plan_from_state( + "SELECT age + 1 AS k FROM person WHERE array_reduce([1.0, 2.0], 0, (age, v) -> age + k, v -> -v) > 0", + &DatabricksDialect, + ParserOptions::default(), + state, + ) + .unwrap(); + let plan = plan.display_indent().to_string(); + assert_contains!(plan, "__DF_WHERE_ALIAS_LAMBDA_0 + person.age"); + + let mut state = MockSessionState::default() + .with_higher_order_function(Arc::new(HigherOrderUDF::new_from_impl( + MockArrayReduce::new(), + ))) + .with_scalar_function(make_array_udf()) + .with_expr_planner(Arc::new(CustomExprPlanner {})); + state.config_options.sql_parser.dialect = + datafusion_common::config::Dialect::Snowflake; + let plan = logical_plan_from_state( + "SELECT age + 1 AS k FROM (SELECT 10 AS age) WHERE array_reduce([1.0, 2.0], 0, (age, v) -> age + k, v -> -v) > 0", + &DatabricksDialect, + ParserOptions::default(), + state, + ) + .unwrap(); + let plan = plan.display_indent().to_string(); + assert_contains!(plan, "__DF_WHERE_ALIAS_LAMBDA_0 + age"); +} + +#[test] +fn snowflake_where_alias_expansion_is_bounded() { + let mut config = datafusion_common::config::ConfigOptions::default(); + config.sql_parser.dialect = datafusion_common::config::Dialect::Snowflake; + + let mut aliases = vec!["1 AS a0".to_string()]; + for index in 1..20 { + aliases.push(format!( + "a{previous} + a{previous} AS a{index}", + previous = index - 1 + )); + } + let sql = format!("SELECT {} WHERE a19 > 0", aliases.join(", ")); + let err = logical_plan_with_config(&sql, config).unwrap_err(); + assert_contains!(err.to_string(), "WHERE alias expansion exceeds"); +} + fn logical_plan_with_config( sql: &str, config_options: datafusion_common::config::ConfigOptions,