From dc7d9d2625d34c7a93eda1fa717dc1e2b66b5046 Mon Sep 17 00:00:00 2001 From: osipovartem Date: Sat, 5 Sep 2026 17:26:03 +0300 Subject: [PATCH] Plan nested window functions for Snowflake --- datafusion/core/tests/sql/select.rs | 28 ++++ datafusion/sql/src/select.rs | 166 ++++++++++++++++++++---- datafusion/sql/tests/sql_integration.rs | 18 +++ 3 files changed, 184 insertions(+), 28 deletions(-) diff --git a/datafusion/core/tests/sql/select.rs b/datafusion/core/tests/sql/select.rs index 96b911e8db130..f287cfd425c4c 100644 --- a/datafusion/core/tests/sql/select.rs +++ b/datafusion/core/tests/sql/select.rs @@ -21,6 +21,34 @@ use super::*; use datafusion_common::{ParamValues, ScalarValue, metadata::ScalarAndMetadata}; use insta::assert_snapshot; +#[tokio::test] +async fn snowflake_nested_window_functions_execute_in_stages() -> Result<()> { + let config = + SessionConfig::new().set_str("datafusion.sql_parser.dialect", "Snowflake"); + let ctx = SessionContext::new_with_config(config); + + let results = ctx + .sql( + "SELECT column1 AS id, \ + SUM(SUM(column1) OVER ()) OVER () AS nested_sum \ + FROM VALUES (1), (2), (3) ORDER BY id", + ) + .await? + .collect() + .await?; + + assert_snapshot!(batches_to_sort_string(&results), @r" + +----+------------+ + | id | nested_sum | + +----+------------+ + | 1 | 18 | + | 2 | 18 | + | 3 | 18 | + +----+------------+ + "); + Ok(()) +} + #[tokio::test] async fn test_list_query_parameters() -> Result<()> { let tmp_dir = TempDir::new()?; diff --git a/datafusion/sql/src/select.rs b/datafusion/sql/src/select.rs index fe601bc08d052..5b5085d4ed353 100644 --- a/datafusion/sql/src/select.rs +++ b/datafusion/sql/src/select.rs @@ -30,6 +30,7 @@ use crate::utils::{ }; 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}; @@ -39,7 +40,8 @@ use datafusion_expr::builder::get_struct_unnested_columns; use datafusion_expr::expr::Unnest as UnnestExpr; use datafusion_expr::expr::{PlannedReplaceSelectItem, WildcardOptions}; use datafusion_expr::expr_rewriter::{ - normalize_col, normalize_col_with_schemas_and_ambiguity_check, normalize_sorts, + NamePreserver, normalize_col, normalize_col_with_schemas_and_ambiguity_check, + normalize_sorts, }; use datafusion_expr::select_expr::SelectExpr; use datafusion_expr::utils::{ @@ -87,6 +89,39 @@ struct RewrittenUnnestExprGroups { expr_groups: Vec>, } +fn contains_nested_window_function(expr: &Expr) -> Result { + let mut found = false; + expr.apply_children(|child| { + child.apply(|nested| { + if matches!(nested, Expr::WindowFunction(_)) { + found = true; + Ok(TreeNodeRecursion::Stop) + } else { + Ok(TreeNodeRecursion::Continue) + } + }) + })?; + Ok(found) +} + +fn find_innermost_window_exprs<'a>( + exprs: impl IntoIterator, +) -> Result> { + let mut window_exprs = Vec::new(); + for expr in exprs { + expr.apply(|nested| { + if matches!(nested, Expr::WindowFunction(_)) + && !contains_nested_window_function(nested)? + && !window_exprs.contains(nested) + { + window_exprs.push(nested.clone()); + } + Ok(TreeNodeRecursion::Continue) + })?; + } + Ok(window_exprs) +} + fn flatten_expr_groups(expr_groups: Vec>) -> Vec { expr_groups.into_iter().flatten().collect() } @@ -351,7 +386,7 @@ impl SqlToRel<'_, S> { plan, select_exprs: mut select_exprs_post_aggr, having_expr: having_expr_post_aggr, - qualify_expr: qualify_expr_post_aggr, + qualify_expr: mut qualify_expr_post_aggr, order_by_exprs: mut order_by_rex, on_exprs: mut on_exprs_post_aggr, } = if !group_by_exprs.is_empty() || !aggr_exprs.is_empty() { @@ -401,11 +436,71 @@ impl SqlToRel<'_, S> { .chain(order_by_rex.iter().map(|s| &s.expr)) .chain(on_exprs_post_aggr.iter()), ); + let qualify_had_window_functions = !find_window_exprs( + select_exprs_post_aggr + .iter() + .chain(qualify_expr_post_aggr.iter()), + ) + .is_empty(); + let plan_nested_windows = self.context_provider.options().sql_parser.dialect + == SqlDialect::Snowflake + && window_func_exprs.iter().try_fold(false, |found, expr| { + Ok::<_, datafusion_common::DataFusionError>( + found || contains_nested_window_function(expr)?, + ) + })?; // Process window functions after aggregation as they can reference // aggregate functions in their body let plan = if window_func_exprs.is_empty() { plan + } else if plan_nested_windows { + // Snowflake permits a window call inside another window call. A logical + // Window node cannot evaluate that shape directly, so materialize each + // innermost level and rebase its parents onto the generated columns. + let mut plan = plan; + loop { + let window_level = find_innermost_window_exprs( + select_exprs_post_aggr + .iter() + .chain(qualify_expr_post_aggr.iter()) + .chain(order_by_rex.iter().map(|sort| &sort.expr)) + .chain(on_exprs_post_aggr.iter()), + )?; + if window_level.is_empty() { + break; + } + + plan = LogicalPlanBuilder::window_plan(plan, window_level.clone())?; + let name_preserver = NamePreserver::new_for_projection(); + select_exprs_post_aggr = select_exprs_post_aggr + .iter() + .map(|expr| { + let saved_name = name_preserver.save(expr); + rebase_expr(expr, &window_level, &plan) + .map(|expr| saved_name.restore(expr)) + }) + .collect::>>()?; + qualify_expr_post_aggr = qualify_expr_post_aggr + .as_ref() + .map(|expr| rebase_expr(expr, &window_level, &plan)) + .transpose()?; + order_by_rex = order_by_rex + .into_iter() + .map(|sort_expr| { + Ok(sort_expr.with_expr(rebase_expr( + &sort_expr.expr, + &window_level, + &plan, + )?)) + }) + .collect::>>()?; + on_exprs_post_aggr = on_exprs_post_aggr + .iter() + .map(|expr| rebase_expr(expr, &window_level, &plan)) + .collect::>>()?; + } + plan } else { let plan = LogicalPlanBuilder::window_plan(plan, window_func_exprs.clone())?; @@ -437,37 +532,52 @@ impl SqlToRel<'_, S> { // Process QUALIFY clause after window functions // QUALIFY filters the results of window functions, similar to how HAVING filters aggregates let plan = if let Some(qualify_expr) = qualify_expr_post_aggr { - // Validate that QUALIFY is used with window functions in SELECT or QUALIFY - let qualify_window_func_exprs = find_window_exprs( - select_exprs_post_aggr - .iter() - .chain(std::iter::once(&qualify_expr)), - ); - if qualify_window_func_exprs.is_empty() { - return plan_err!( - "QUALIFY clause requires window functions in the SELECT list or QUALIFY clause" + if plan_nested_windows { + if !qualify_had_window_functions { + return plan_err!( + "QUALIFY clause requires window functions in the SELECT list or QUALIFY clause" + ); + } + self.validate_schema_satisfies_exprs( + plan.schema(), + std::slice::from_ref(&qualify_expr), + )?; + LogicalPlanBuilder::from(plan) + .filter(qualify_expr)? + .build()? + } else { + // Validate that QUALIFY is used with window functions in SELECT or QUALIFY + let qualify_window_func_exprs = find_window_exprs( + select_exprs_post_aggr + .iter() + .chain(std::iter::once(&qualify_expr)), ); - } + if qualify_window_func_exprs.is_empty() { + return plan_err!( + "QUALIFY clause requires window functions in the SELECT list or QUALIFY clause" + ); + } - // now attempt to resolve columns and replace with fully-qualified columns - let windows_projection_exprs = window_func_exprs - .iter() - .map(|expr| resolve_columns(expr, &plan)) - .collect::>>()?; + // now attempt to resolve columns and replace with fully-qualified columns + let windows_projection_exprs = window_func_exprs + .iter() + .map(|expr| resolve_columns(expr, &plan)) + .collect::>>()?; - // Rewrite the qualify expression to reference columns from the window plan - let qualify_expr_post_window = - rebase_expr(&qualify_expr, &windows_projection_exprs, &plan)?; + // Rewrite the qualify expression to reference columns from the window plan + let qualify_expr_post_window = + rebase_expr(&qualify_expr, &windows_projection_exprs, &plan)?; - // Validate that the qualify expression can be resolved from the window plan schema - self.validate_schema_satisfies_exprs( - plan.schema(), - std::slice::from_ref(&qualify_expr_post_window), - )?; + // Validate that the qualify expression can be resolved from the window plan schema + self.validate_schema_satisfies_exprs( + plan.schema(), + std::slice::from_ref(&qualify_expr_post_window), + )?; - LogicalPlanBuilder::from(plan) - .filter(qualify_expr_post_window)? - .build()? + LogicalPlanBuilder::from(plan) + .filter(qualify_expr_post_window)? + .build()? + } } else { plan }; diff --git a/datafusion/sql/tests/sql_integration.rs b/datafusion/sql/tests/sql_integration.rs index 85c1eda25d8ee..bbda33dac95fd 100644 --- a/datafusion/sql/tests/sql_integration.rs +++ b/datafusion/sql/tests/sql_integration.rs @@ -2014,6 +2014,24 @@ fn select_nested_window_function() { ); } +#[test] +fn select_nested_window_function_snowflake() { + let mut config_options = datafusion_common::config::ConfigOptions::new(); + config_options.sql_parser.dialect = datafusion_common::config::Dialect::Snowflake; + + let plan = logical_plan_with_config( + "SELECT sum(sum(age) OVER ()) OVER () FROM person", + config_options, + ) + .unwrap(); + assert_snapshot!(plan, @r" + Projection: sum(sum(person.age) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING + WindowAggr: windowExpr=[[sum(sum(person.age) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING]] + WindowAggr: windowExpr=[[sum(person.age) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING]] + TableScan: person + "); +} + #[test] fn select_aggregate_inside_window_function() { // an aggregate as the argument of a window function is legal: the window