From 4faee11eef76fea2d9bef0d13f6300160e663bae Mon Sep 17 00:00:00 2001 From: osipovartem Date: Sun, 6 Sep 2026 15:19:08 +0300 Subject: [PATCH] Revert "Plan nested window functions for Snowflake (#67)" This reverts commit 0676b9cbd9fd5fdc1cd37ef30e02ff20b16af3ad. --- datafusion/core/tests/sql/select.rs | 28 ---- datafusion/sql/src/select.rs | 166 ++++-------------------- datafusion/sql/tests/sql_integration.rs | 14 +- 3 files changed, 34 insertions(+), 174 deletions(-) diff --git a/datafusion/core/tests/sql/select.rs b/datafusion/core/tests/sql/select.rs index f287cfd425c4c..96b911e8db130 100644 --- a/datafusion/core/tests/sql/select.rs +++ b/datafusion/core/tests/sql/select.rs @@ -21,34 +21,6 @@ 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 5b5085d4ed353..fe601bc08d052 100644 --- a/datafusion/sql/src/select.rs +++ b/datafusion/sql/src/select.rs @@ -30,7 +30,6 @@ 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}; @@ -40,8 +39,7 @@ 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::{ - NamePreserver, normalize_col, normalize_col_with_schemas_and_ambiguity_check, - normalize_sorts, + normalize_col, normalize_col_with_schemas_and_ambiguity_check, normalize_sorts, }; use datafusion_expr::select_expr::SelectExpr; use datafusion_expr::utils::{ @@ -89,39 +87,6 @@ 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() } @@ -386,7 +351,7 @@ impl SqlToRel<'_, S> { plan, select_exprs: mut select_exprs_post_aggr, having_expr: having_expr_post_aggr, - qualify_expr: mut qualify_expr_post_aggr, + qualify_expr: 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() { @@ -436,71 +401,11 @@ 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())?; @@ -532,52 +437,37 @@ 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 { - 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)), + // 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 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 d387f6decaa8b..18ff4c26b7acb 100644 --- a/datafusion/sql/tests/sql_integration.rs +++ b/datafusion/sql/tests/sql_integration.rs @@ -2035,17 +2035,15 @@ 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( + let err = 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 - "); + .expect_err("Snowflake rejects nested window function calls"); + assert_snapshot!( + err.strip_backtrace(), + @"Error during planning: Window function calls cannot be nested: 'sum(person.age) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING' is nested inside 'sum(sum(person.age) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING'" + ); } #[test]