Repository navigation
fix: do not extract common subexpressions from the arguments of an aggregate with a FILTER - #26184
Conversation
…gregate with a FILTER
Codecov Report❌ Patch coverage is
Additional details and impacted files☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
@kosiew would you mind reviewing this change? Thanks! |
| // follow the same order. | ||
| node.apply_children(|child| { | ||
| self.conditional = | ||
| conditional.iter().any(|c| std::ptr::eq(*c, child)); |
There was a problem hiding this comment.
could we use a set of child pointers here? Scanning conditional for every child makes wide CASE traversal quadratic. In an optimized CSE probe, 10k arms went from ~1.4 ms to ~39 ms, and 40k from ~7.4 ms to ~616 ms (base vs PR, CSE alone rather than whole-query planning). A set would preserve tree order without the repeated scans.
kosiew
left a comment
There was a problem hiding this comment.
Thanks for working on this. The fix looks good to me. Filtered aggregate arguments now stay conditional unless they are also evaluated unconditionally, and the traversal change keeps CSE identifiers aligned with tree order. The new regression tests cover the original failure and the important contrasting cases.
I have one minor documentation suggestion about filtered window arguments, but it is not blocking.
| .collect(), | ||
| )), | ||
|
|
||
| // The same applies to the arguments of a window function with a `FILTER` |
There was a problem hiding this comment.
Could we clarify that filtered window arguments are treated conservatively by CSE, rather than actually being evaluated only for rows passing the filter? AggregateWindowExpr::get_result_column currently evaluates arguments before applying the filter. Please update this comment and the test comment around line 2139 to reflect that distinction.
|
Thanks for the review @kosiew ! |
apache#26356) ## Which issue does this PR close? No issue filed. The reproduction is below. ## Rationale for this change `EliminateGroupByConstant` re-plans the original group expressions over the new `Aggregate` with `LogicalPlanBuilder::project`. Column normalization falls back to the input schema of the aggregate. Thus an unqualified group expression, such as an extracted leaf `__datafusion_extracted_1`, becomes `m.__datafusion_extracted_1` when the input is a `SubqueryAlias: m`. The new aggregate does not output that name, and the rule fails: ``` Optimizer rule 'eliminate_group_by_constant' failed caused by Schema error: No field named m.__datafusion_extracted_1. ``` The input plan is valid, and the rule turns it into an invalid one. We reach it when we optimize a plan, bind its parameters with `replace_params_with_values`, and optimize again. The bound value makes a group key redundant only after `ExtractLeafExpressions` has rewritten the plan. In a single `optimize` call, `EliminateGroupByConstant` runs before `ExtractLeafExpressions` in each pass, so I could not reproduce it there. Since apache#26184, the arguments of filtered aggregates no longer go into a common-expression projection between the aggregate and the alias, so the aggregate can sit directly on the alias and the problem shows. ```rust let ctx = SessionContext::new(); ctx.sql("CREATE TABLE t AS SELECT * FROM (VALUES (TIMESTAMP '2026-01-01T00:00:00', named_struct('a', 'x'), 1)) AS v(ts, s, x)") .await?.collect().await?; let scan = ctx.table("t").await?.into_unoptimized_plan(); let plan = LogicalPlanBuilder::from(scan) .alias("m")? .aggregate( vec![ date_bin( cast(placeholder("$1"), DataType::Interval(IntervalUnit::MonthDayNano)), col("ts"), lit(ScalarValue::TimestampNanosecond(Some(0), None)), ), get_field(col("s"), "a").alias("c"), col("ts"), ], vec![ sum(col("x") + lit(1)).filter(col("x").gt(lit(0))).build()?, max(col("x") + lit(1)).filter(col("x").lt(lit(5))).build()?, ], )? .build()?; let state = ctx.state(); let plan = state.optimize(&plan)?; let plan = plan.replace_params_with_values(&ParamValues::List(vec![ ScalarValue::Utf8(Some("1 minute".into())).into(), ]))?; state.optimize(&plan)?; // fails with the error above ``` ## What changes are included in this PR? - The rule builds the projection with `Projection::try_new` and does not re-plan the group expressions. A kept key or an aggregate becomes a column of the schema of the new aggregate. A redundant key keeps its expression, which reads only bare GROUP BY columns that the new aggregate outputs unchanged. An item gets an alias only when its name differs from the original output field, so existing plans do not change. ## What is the testing strategy for this PR? - New unit test `test_eliminate_with_unqualified_reference_to_qualified_input`. It fails without the change, with the same schema error. - `cargo test -p datafusion-optimizer` and `cargo test -p datafusion-sqllogictest --test sqllogictests` pass. No existing plan changes. ## Are there any user-facing changes? No API changes. A plan that failed in this rule now optimizes. --------- Co-authored-by: Copilot Autofix powered by AI <[email protected]> Co-authored-by: Claude Opus 5.5 <[email protected]>
Which issue does this PR close?
FILTERfor rows that the filter rejects #26183.Rationale for this change
An aggregate with a
FILTERclause evaluates its arguments only for the rows that pass the filter, so a fallible argument does not fail on a rejected row.When two aggregates have the same argument,
CommonSubexprEliminatemoves that argument into a projection below the aggregate.The projection evaluates it for every row, and the query fails on a row that the filter rejects:
The rule already has a model for this: a conditionally evaluated subexpression (a
CASEbranch, the right side ofAND/OR, a lazy function argument) is extracted only if it is also surely evaluated somewhere else.The arguments of a filtered aggregate are conditionally evaluated in the same sense, but the rule did not know that.
What changes are included in this PR?
ExprCSEController::conditional_childrennow reports the arguments and theORDER BYexpressions of anExpr::AggregateFunctionwith aFILTERas conditionally evaluated.The filter itself stays a normal child, because the aggregate evaluates it for every row.
A filter that several aggregates share is still extracted.
Expr::WindowFunctionwith aFILTERas conditionally evaluated.The
PARTITION BYexpressions, theORDER BYexpressions and the filter stay normal children.The window executor does not yet evaluate the filter first (Window aggregate FILTER evaluates arguments for rejected rows #24507, fix: evaluate window aggregate arguments after FILTER #24508), so this part has no visible effect today.
It keeps the rule correct when that fix merges.
CSEVisitor::f_downindatafusion/common/src/cse.rsnow visits the children of a node with conditional children in tree order.Before, it visited all normal children and then all conditional children.
CSERewriterreadsid_arrayin pre-order, so that was correct only when every normal child came before every conditional child in the tree.The children of an aggregate are
args,filter,order_by, so the normal child is in the middle.The
CSEControllertrait is unchanged.The rule keeps every extraction that is safe:
sum(f(x)) FILTER (WHERE p), count(f(x)) FILTER (WHERE p)f(x)andpextractedpextractedsum(f(x)) FILTER (WHERE p), count(f(x))f(x)extractedf(x)extracted (the second aggregate evaluates it for every row)sum(f(x)) FILTER (WHERE f(x) > 0)f(x)extractedf(x)extracted (the filter evaluates it for every row)sum(f(x)) FILTER (WHERE p) AS a, sum(f(x)) FILTER (WHERE p) AS bThe cost is the first row: when an expression occurs only under filtered aggregates, each aggregate now evaluates it, on the rows that pass its filter.
An alternative keeps that extraction when every aggregate has the same filter, with a guard such as
CASE WHEN p THEN f(x) ENDin the projection.I did not do that here because it is a larger change with its own plan shape, and it can be a follow-up.
What is the testing strategy for this PR?
common_subexpr_eliminate.rs:test_aggregate_filter_arguments(the four rows of the table above),test_aggregate_filter_order_byandtest_window_filter_arguments.cse.slt: the reproduction above, with and withoutGROUP BY, theEXPLAINthat shows the cast stays in the aggregate, and the case where an unfiltered aggregate shares the argument.aggregates_simplify.slt: one expected plan changes.SUM(column1 + 1) FILTER (WHERE column1 > 1), SUM(column1 + 1) FILTER (WHERE column1 > 0)no longer extractscolumn1 + 1.Are there any user-facing changes?
Queries that failed with an error from a row that an aggregate
FILTERrejects now succeed.Plans for queries where a subexpression occurs only in the arguments of filtered aggregates no longer have that subexpression in a
__common_exprprojection.There are no API changes.