diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index 2ecb12c30afad..1f32d9c6da445 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -2900,6 +2900,48 @@ mod tests { Ok(()) } + #[test] + fn plan_builder_aggregate_rejects_nested_aggregates() -> Result<()> { + // https://github.com/apache/datafusion/issues/23812 + let err = table_scan( + Some("employee_csv"), + &employee_schema(), + Some(vec![0, 3, 4]), + )? + .aggregate(vec![col("id")], vec![sum(sum(col("salary")))]) + .expect_err("nested aggregates should be rejected"); + + assert_snapshot!( + err.strip_backtrace(), + @"Error during planning: Aggregate function calls cannot be nested: 'sum(employee_csv.salary)' is nested inside 'sum(sum(employee_csv.salary))'" + ); + + Ok(()) + } + + #[test] + fn plan_builder_window_rejects_nested_window_functions() -> Result<()> { + // https://github.com/apache/datafusion/issues/23812 + let sum_over = |arg| { + Expr::from(expr::WindowFunction::new( + crate::WindowFunctionDefinition::AggregateUDF( + crate::test::function_stub::sum_udaf(), + ), + vec![arg], + )) + }; + let err = table_scan(Some("employee_csv"), &employee_schema(), Some(vec![4]))? + .window(vec![sum_over(sum_over(col("salary")))]) + .expect_err("nested window functions should be rejected"); + + assert_snapshot!( + err.strip_backtrace(), + @"Error during planning: Window function calls cannot be nested: 'sum(employee_csv.salary) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING' is nested inside 'sum(sum(employee_csv.salary) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING'" + ); + + Ok(()) + } + #[test] fn test_join_metadata() -> Result<()> { let left_schema = DFSchema::new_with_metadata( diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index 9cfab21a0395e..9ac27b46a78e6 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -41,8 +41,9 @@ use crate::logical_plan::display::{GraphvizVisitor, IndentVisitor}; use crate::logical_plan::extension::UserDefinedLogicalNode; use crate::logical_plan::{DmlStatement, Statement}; use crate::utils::{ - enumerate_grouping_sets, exprlist_to_fields, find_out_reference_exprs, - grouping_set_expr_count, grouping_set_to_exprlist, merge_schema, split_conjunction, + check_aggregate_and_window_nesting, enumerate_grouping_sets, exprlist_to_fields, + find_out_reference_exprs, grouping_set_expr_count, grouping_set_to_exprlist, + merge_schema, split_conjunction, }; use crate::{ BinaryExpr, CreateMemoryTable, CreateView, Execute, Expr, ExprSchemable, GroupingSet, @@ -2779,6 +2780,11 @@ pub struct Window { impl Window { /// Create a new window operator. pub fn try_new(window_expr: Vec, input: Arc) -> Result { + // Reject e.g. `sum(sum(x) OVER ()) OVER ()` here rather than letting it + // reach physical planning, which has no equivalent for a nested window + // function. + check_aggregate_and_window_nesting(window_expr.iter())?; + let fields: Vec<(Option, Arc)> = input .schema() .iter() @@ -3892,6 +3898,10 @@ impl Aggregate { group_expr: Vec, aggr_expr: Vec, ) -> Result { + // Reject e.g. `sum(sum(x))` here rather than letting it reach physical + // planning, which has no equivalent for a nested aggregate. + check_aggregate_and_window_nesting(group_expr.iter().chain(aggr_expr.iter()))?; + let group_expr = enumerate_grouping_sets(group_expr)?; let is_grouping_set = matches!(group_expr.as_slice(), [Expr::GroupingSet(_)]); diff --git a/datafusion/expr/src/utils.rs b/datafusion/expr/src/utils.rs index 22abb454d4e6b..7f79c5cf18c4a 100644 --- a/datafusion/expr/src/utils.rs +++ b/datafusion/expr/src/utils.rs @@ -34,8 +34,8 @@ use datafusion_common::tree_node::{ }; use datafusion_common::utils::get_at_indices; use datafusion_common::{ - Column, DFSchema, DFSchemaRef, HashMap, Result, TableReference, internal_err, - plan_err, + Column, DFSchema, DFSchemaRef, DataFusionError, Diagnostic, HashMap, Result, Span, + TableReference, internal_err, plan_datafusion_err, plan_err, }; #[cfg(not(feature = "sql"))] @@ -652,6 +652,106 @@ pub fn find_aggregate_exprs<'a>(exprs: impl IntoIterator) -> Ve }) } +/// Returns an error if any of `exprs` nests aggregate or window function calls +/// in a way that has no physical equivalent: an aggregate call may not contain +/// another aggregate call (`sum(sum(x))`) or a window call +/// (`sum(sum(x) OVER ())`), and a window call may not contain another window +/// call (`sum(sum(x) OVER ()) OVER ()`). The reverse nesting, an aggregate used +/// as the argument of a window call (`sum(sum(x)) OVER ()`), is legal: there the +/// aggregate is evaluated by the `Aggregate` node and the window function is +/// evaluated on top of its result. +/// +/// Such expressions are not valid SQL either, so they are rejected while the +/// logical plan is built rather than failing later with an error that does not +/// point back at the original SQL. +/// +/// [`Aggregate::try_new`] and [`Window::try_new`] call this, so the SQL planner +/// and the `DataFrame`/`LogicalPlanBuilder` paths are checked without callers +/// invoking it directly. The lower-level `try_new_with_schema` constructors and +/// building a `Window` from its public fields bypass the check, so a caller +/// that constructs those nodes by hand should call this itself. +/// +/// [`Aggregate::try_new`]: crate::logical_plan::Aggregate::try_new +/// [`Window::try_new`]: crate::logical_plan::Window::try_new +pub(crate) fn check_aggregate_and_window_nesting<'a>( + exprs: impl IntoIterator, +) -> Result<()> { + for expr in exprs { + expr.apply(|outer| { + if !matches!(outer, Expr::AggregateFunction(_) | Expr::WindowFunction(_)) { + return Ok(TreeNodeRecursion::Continue); + } + + // Look for an illegally nested call in the arguments, `FILTER`, + // `ORDER BY` and `PARTITION BY` of this call + let mut err = None; + outer.apply_children(|child| { + child.apply(|inner| { + err = illegal_nesting_err(outer, inner); + if err.is_some() { + Ok(TreeNodeRecursion::Stop) + } else { + Ok(TreeNodeRecursion::Continue) + } + }) + })?; + + match err { + Some(err) => Err(err), + None => Ok(TreeNodeRecursion::Continue), + } + })?; + } + Ok(()) +} + +/// The planning error for a call to `inner` nested inside a call to `outer`, or +/// `None` if that nesting is legal. +fn illegal_nesting_err(outer: &Expr, inner: &Expr) -> Option { + // Messages follow PostgreSQL, which rejects the same three cases + let (message, help) = match (outer, inner) { + (Expr::AggregateFunction(_), Expr::AggregateFunction(_)) => ( + "Aggregate function calls cannot be nested", + format!("Compute '{inner}' in an inner query and aggregate its result"), + ), + (Expr::AggregateFunction(_), Expr::WindowFunction(_)) => ( + "Aggregate function calls cannot contain window function calls", + format!("Compute '{inner}' in an inner query and aggregate its result"), + ), + (Expr::WindowFunction(_), Expr::WindowFunction(_)) => ( + "Window function calls cannot be nested", + format!("Compute '{inner}' in an inner query and use its result here"), + ), + // Anything else, including an aggregate inside a window call + _ => return None, + }; + + Some( + plan_datafusion_err!("{message}: '{inner}' is nested inside '{outer}'") + .with_diagnostic( + Diagnostic::new_error(message, first_span(inner)).with_help(help, None), + ), + ) +} + +/// Best effort source location for `expr`: the first [`Span`] found in its +/// subtree. Only some expressions (currently columns) carry spans, so pointing +/// at e.g. the column of `sum(x)` is the closest we can get to the location of +/// the whole expression. +fn first_span(expr: &Expr) -> Option { + let mut span = None; + expr.apply(|e| { + span = e.spans().and_then(|spans| spans.first()); + if span.is_some() { + Ok(TreeNodeRecursion::Stop) + } else { + Ok(TreeNodeRecursion::Continue) + } + }) + .ok()?; + span +} + /// Collect all deeply nested `Expr::WindowFunction`. They are returned in order of occurrence /// (depth first), with duplicates omitted. pub fn find_window_exprs<'a>(exprs: impl IntoIterator) -> Vec { @@ -1917,4 +2017,91 @@ mod tests { substr(string: String, start_pos: Int64, length: Int64) "); } + + /// `sum() OVER ()` + fn sum_over(args: Vec) -> Expr { + Expr::from(WindowFunction::new( + WindowFunctionDefinition::AggregateUDF(sum_udaf()), + args, + )) + } + + #[test] + fn test_check_aggregate_and_window_nesting_ok() -> Result<()> { + use crate::test::function_stub::{count, sum}; + + let exprs = [ + // a plain aggregate, and one wrapped in a scalar expression + sum(col("a")), + count(col("a")) + lit(1), + // a window function over a column, and over an aggregate + sum_over(vec![col("a")]), + sum_over(vec![sum(col("a"))]), + ]; + + check_aggregate_and_window_nesting(exprs.iter())?; + Ok(()) + } + + #[test] + fn test_check_aggregate_and_window_nesting_err() { + use crate::test::function_stub::{count, sum}; + use insta::assert_snapshot; + + // an aggregate directly inside an aggregate + let err = check_aggregate_and_window_nesting([&sum(sum(col("a")))]).unwrap_err(); + assert_snapshot!( + err.strip_backtrace(), + @"Error during planning: Aggregate function calls cannot be nested: 'sum(a)' is nested inside 'sum(sum(a))'" + ); + + // nested below another expression in the arguments + let err = check_aggregate_and_window_nesting([&sum(col("a") + count(col("b")))]) + .unwrap_err(); + assert_snapshot!( + err.strip_backtrace(), + @"Error during planning: Aggregate function calls cannot be nested: 'COUNT(b)' is nested inside 'sum(a + COUNT(b))'" + ); + + // nested in the FILTER of an aggregate + let filtered = sum(col("a")) + .filter(sum(col("b")).gt(lit(0))) + .build() + .unwrap(); + let err = check_aggregate_and_window_nesting([&filtered]).unwrap_err(); + assert_snapshot!( + err.strip_backtrace(), + @"Error during planning: Aggregate function calls cannot be nested: 'sum(b)' is nested inside 'sum(a) FILTER (WHERE sum(b) > Int32(0))'" + ); + + // nested in the ORDER BY of an aggregate + let ordered = sum(col("a")) + .order_by(vec![Sort::new(sum(col("b")), true, false)]) + .build() + .unwrap(); + let err = check_aggregate_and_window_nesting([&ordered]).unwrap_err(); + assert_snapshot!( + err.strip_backtrace(), + @"Error during planning: Aggregate function calls cannot be nested: 'sum(b)' is nested inside 'sum(a) ORDER BY [sum(b) ASC NULLS LAST]'" + ); + + // a window function inside an aggregate + let err = check_aggregate_and_window_nesting([&sum(sum_over(vec![col("a")]))]) + .unwrap_err(); + assert_snapshot!( + err.strip_backtrace(), + @"Error during planning: Aggregate function calls cannot contain window function calls: 'sum(a) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING' is nested inside 'sum(sum(a) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING)'" + ); + + // a window function inside a window function + let err = + check_aggregate_and_window_nesting([&sum_over(vec![sum_over(vec![col( + "a", + )])])]) + .unwrap_err(); + assert_snapshot!( + err.strip_backtrace(), + @"Error during planning: Window function calls cannot be nested: 'sum(a) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING' is nested inside 'sum(sum(a) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING'" + ); + } } diff --git a/datafusion/sql/tests/cases/diagnostic.rs b/datafusion/sql/tests/cases/diagnostic.rs index df46a48d88579..1f2cefdec0629 100644 --- a/datafusion/sql/tests/cases/diagnostic.rs +++ b/datafusion/sql/tests/cases/diagnostic.rs @@ -16,6 +16,7 @@ // under the License. use datafusion_functions::string; +use datafusion_functions_aggregate::sum::sum_udaf; use insta::assert_snapshot; use std::{collections::HashMap, ops::ControlFlow, sync::Arc}; @@ -44,7 +45,8 @@ fn do_query(sql: &'static str) -> Diagnostic { ..ParserOptions::default() }; let state = MockSessionState::default() - .with_scalar_function(Arc::new(string::concat().as_ref().clone())); + .with_scalar_function(Arc::new(string::concat().as_ref().clone())) + .with_aggregate_function(sum_udaf()); let context = MockContextProvider { state }; let sql_to_rel = SqlToRel::new_with_options(&context, options); match sql_to_rel.statement_to_plan(statement) { @@ -671,3 +673,40 @@ fn test_multiple_null_comparison_warnings() -> Result<()> { ); Ok(()) } + +#[test] +fn test_nested_aggregate() -> Result<()> { + let query = "SELECT sum(sum(/*a*/age/*a*/)) FROM person"; + let spans = get_spans(query); + let diag = do_query(query); + assert_snapshot!(diag.message, @"Aggregate function calls cannot be nested"); + assert_eq!(diag.span, Some(spans["a"])); + assert_snapshot!( + diag.helps[0].message, + @"Compute 'sum(person.age)' in an inner query and aggregate its result" + ); + Ok(()) +} + +#[test] +fn test_window_function_inside_aggregate() -> Result<()> { + let query = "SELECT sum(sum(/*a*/age/*a*/) OVER ()) FROM person"; + let spans = get_spans(query); + let diag = do_query(query); + assert_snapshot!( + diag.message, + @"Aggregate function calls cannot contain window function calls" + ); + assert_eq!(diag.span, Some(spans["a"])); + Ok(()) +} + +#[test] +fn test_nested_window_function() -> Result<()> { + let query = "SELECT sum(sum(/*a*/age/*a*/) OVER ()) OVER () FROM person"; + let spans = get_spans(query); + let diag = do_query(query); + assert_snapshot!(diag.message, @"Window function calls cannot be nested"); + assert_eq!(diag.span, Some(spans["a"])); + Ok(()) +} diff --git a/datafusion/sql/tests/sql_integration.rs b/datafusion/sql/tests/sql_integration.rs index 54480a4224992..a4bf0db910774 100644 --- a/datafusion/sql/tests/sql_integration.rs +++ b/datafusion/sql/tests/sql_integration.rs @@ -1771,6 +1771,81 @@ fn select_simple_aggregate_with_groupby_position_out_of_range() { ); } +#[test] +fn select_nested_aggregate() { + // https://github.com/apache/datafusion/issues/23812 + let err = logical_plan("SELECT sum(sum(age)) FROM person") + .expect_err("query should have failed"); + assert_snapshot!( + err.strip_backtrace(), + @"Error during planning: Aggregate function calls cannot be nested: 'sum(person.age)' is nested inside 'sum(sum(person.age))'" + ); + + let err = logical_plan("SELECT state, sum(count(age)) FROM person GROUP BY state") + .expect_err("query should have failed"); + assert_snapshot!( + err.strip_backtrace(), + @"Error during planning: Aggregate function calls cannot be nested: 'count(person.age)' is nested inside 'sum(count(person.age))'" + ); + + let err = + logical_plan("SELECT state FROM person GROUP BY state HAVING sum(sum(age)) > 0") + .expect_err("query should have failed"); + assert_snapshot!( + err.strip_backtrace(), + @"Error during planning: Aggregate function calls cannot be nested: 'sum(person.age)' is nested inside 'sum(sum(person.age))'" + ); +} + +#[test] +fn select_window_function_inside_aggregate() { + // https://github.com/apache/datafusion/issues/23812 + let err = logical_plan("SELECT sum(sum(age) OVER ()) FROM person") + .expect_err("query should have failed"); + assert_snapshot!( + err.strip_backtrace(), + @"Error during planning: Aggregate function calls cannot contain window function calls: 'sum(person.age) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING' is nested inside 'sum(sum(person.age) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING)'" + ); +} + +#[test] +fn select_nested_window_function() { + // https://github.com/apache/datafusion/issues/23812 + let err = logical_plan("SELECT sum(sum(age) OVER ()) OVER () FROM person") + .expect_err("query should have failed"); + 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'" + ); + + let err = logical_plan( + "SELECT rank() OVER (ORDER BY rank() OVER (ORDER BY age)) FROM person", + ) + .expect_err("query should have failed"); + assert_snapshot!( + err.strip_backtrace(), + @"Error during planning: Window function calls cannot be nested: 'rank() ORDER BY [person.age ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW' is nested inside 'rank() ORDER BY [rank() ORDER BY [person.age ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW'" + ); +} + +#[test] +fn select_aggregate_inside_window_function() { + // an aggregate as the argument of a window function is legal: the window + // function is evaluated on top of the aggregate + let plan = + logical_plan("SELECT state, sum(sum(age)) OVER () FROM person GROUP BY state") + .unwrap(); + assert_snapshot!( + plan, + @r" + Projection: person.state, sum(sum(person.age)) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING + WindowAggr: windowExpr=[[sum(sum(person.age)) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING]] + Aggregate: groupBy=[[person.state]], aggr=[[sum(person.age)]] + TableScan: person + " + ); +} + #[test] fn select_simple_aggregate_with_groupby_can_use_alias() { let plan = diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt index cbb9c5d0317dc..d307ddc066460 100644 --- a/datafusion/sqllogictest/test_files/aggregate.slt +++ b/datafusion/sqllogictest/test_files/aggregate.slt @@ -9607,3 +9607,73 @@ SET datafusion.execution.target_partitions = 4; statement ok DROP TABLE hits_raw; + +# Nested aggregate function calls are rejected during planning +# issue: https://github.com/apache/datafusion/issues/23812 +statement ok +CREATE TABLE nested_agg_t AS VALUES (1, 10.0), (1, 20.0), (2, 30.0); + +statement error DataFusion error: Error during planning: Aggregate function calls cannot be nested: 'sum\(nested_agg_t\.column2\)' is nested inside 'sum\(sum\(nested_agg_t\.column2\)\)' +SELECT column1, sum(sum(column2)) FROM nested_agg_t GROUP BY column1; + +statement error DataFusion error: Error during planning: Aggregate function calls cannot be nested: 'count\(nested_agg_t\.column2\)' is nested inside 'sum\(count\(nested_agg_t\.column2\)\)' +SELECT column1, sum(count(column2)) FROM nested_agg_t GROUP BY column1; + +statement error DataFusion error: Error during planning: Aggregate function calls cannot be nested +SELECT sum(sum(column2)) FROM nested_agg_t; + +statement error DataFusion error: Error during planning: Aggregate function calls cannot be nested +SELECT column1 FROM nested_agg_t GROUP BY column1 HAVING sum(sum(column2)) > 0; + +statement error DataFusion error: Error during planning: Aggregate function calls cannot be nested +SELECT column1, sum(column2 + sum(column2)) FROM nested_agg_t GROUP BY column1; + +statement error DataFusion error: Error during planning: Aggregate function calls cannot be nested +SELECT sum(column2) FILTER (WHERE sum(column2) > 0) FROM nested_agg_t; + +statement error DataFusion error: Error during planning: Aggregate function calls cannot be nested +SELECT array_agg(column2 ORDER BY sum(column2)) FROM nested_agg_t; + +# A window function nested inside an aggregate is rejected as well +statement error DataFusion error: Error during planning: Aggregate function calls cannot contain window function calls +SELECT sum(sum(column2) OVER ()) FROM nested_agg_t; + +# ... as are nested window functions +statement error DataFusion error: Error during planning: Window function calls cannot be nested +SELECT sum(sum(column2) OVER ()) OVER () FROM nested_agg_t; + +statement error DataFusion error: Error during planning: Window function calls cannot be nested +SELECT row_number() OVER (ORDER BY row_number() OVER ()) FROM nested_agg_t; + +# ... including a window call nested in `PARTITION BY` +statement error DataFusion error: Error during planning: Window function calls cannot be nested +SELECT row_number() OVER (PARTITION BY row_number() OVER ()) FROM nested_agg_t; + +# A scalar function applied to an aggregate is legal +query IR +SELECT column1, abs(sum(column2)) FROM nested_agg_t GROUP BY column1 ORDER BY column1; +---- +1 30 +2 30 + +# A window function applied to an aggregate is legal +query IR +SELECT column1, sum(sum(column2)) OVER () FROM nested_agg_t GROUP BY column1 ORDER BY column1; +---- +1 60 +2 60 + +# An aggregate over the result of an aggregate computed in a subquery is legal +query R +SELECT sum(s) FROM (SELECT sum(column2) AS s FROM nested_agg_t GROUP BY column1); +---- +60 + +# An aggregate over the result of a window function computed in a subquery is legal +query R +SELECT sum(s) FROM (SELECT sum(column2) OVER () AS s FROM nested_agg_t); +---- +180 + +statement ok +DROP TABLE nested_agg_t;