Skip to content

fix: do not extract common subexpressions from the arguments of an aggregate with a FILTER - #26184

Merged
alamb merged 3 commits into
apache:mainfrom
pydantic:cse-aggregate-filter-args
Oct 12, 2026
Merged

alamb merged 3 commits into
apache:mainfrom
pydantic:cse-aggregate-filter-args

Conversation

@adriangb

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Rationale for this change

An aggregate with a FILTER clause 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, CommonSubexprEliminate moves 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:

CREATE TABLE t AS
SELECT * FROM (VALUES (CAST(-950 AS SMALLINT), 0), (CAST(5 AS SMALLINT), 1)) AS v(x, y);

-- Succeeds.
SELECT avg(arrow_cast(x, 'UInt64')) FILTER (WHERE y > 0) AS a FROM t;

-- Fails: Arrow error: Cast error: Can't cast value -950 to type UInt64
SELECT
    sum(arrow_cast(x, 'UInt64')) FILTER (WHERE y > 0) AS s,
    count(arrow_cast(x, 'UInt64')) FILTER (WHERE y > 0) AS c
FROM t;

The rule already has a model for this: a conditionally evaluated subexpression (a CASE branch, the right side of AND / 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?

  1. ExprCSEController::conditional_children now reports the arguments and the ORDER BY expressions of an Expr::AggregateFunction with a FILTER as 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.
  2. The same function reports the arguments of an Expr::WindowFunction with a FILTER as conditionally evaluated.
    The PARTITION BY expressions, the ORDER BY expressions 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.
  3. CSEVisitor::f_down in datafusion/common/src/cse.rs now visits the children of a node with conditional children in tree order.
    Before, it visited all normal children and then all conditional children.
    CSERewriter reads id_array in 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 CSEController trait is unchanged.

The rule keeps every extraction that is safe:

Query shape Before After
sum(f(x)) FILTER (WHERE p), count(f(x)) FILTER (WHERE p) f(x) and p extracted only p extracted
sum(f(x)) FILTER (WHERE p), count(f(x)) f(x) extracted f(x) extracted (the second aggregate evaluates it for every row)
sum(f(x)) FILTER (WHERE f(x) > 0) f(x) extracted f(x) extracted (the filter evaluates it for every row)
sum(f(x)) FILTER (WHERE p) AS a, sum(f(x)) FILTER (WHERE p) AS b one aggregate one aggregate

The 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) END in 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?

  • Unit tests in common_subexpr_eliminate.rs: test_aggregate_filter_arguments (the four rows of the table above), test_aggregate_filter_order_by and test_window_filter_arguments.
  • cse.slt: the reproduction above, with and without GROUP BY, the EXPLAIN that 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 extracts column1 + 1.

Are there any user-facing changes?

Queries that failed with an error from a row that an aggregate FILTER rejects now succeed.
Plans for queries where a subexpression occurs only in the arguments of filtered aggregates no longer have that subexpression in a __common_expr projection.
There are no API changes.

@github-actions github-actions Bot added optimizer Optimizer rules sqllogictest SQL Logic Tests (.slt) common Related to common crate labels Oct 10, 2026
@codecov-commenter

codecov-commenter commented Oct 10, 2026 •

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 80.79470% with 29 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.82%. Comparing base (334f293) to head (b3004e5).
⚠️ Report is 10 commits behind head on main.

Files with missing lines Patch % Lines
...tafusion/optimizer/src/common_subexpr_eliminate.rs 80.85% 9 Missing and 18 partials ⚠️
datafusion/common/src/cse.rs 80.00% 0 Missing and 2 partials ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main   #26184      +/-   ##
==========================================
+ Coverage   82.79%   82.82%   +0.03%     
==========================================
  Files        1149     1149              
  Lines      451645   452487     +842     
  Branches   451645   452487     +842     
==========================================
+ Hits       373940   374788     +848     
+ Misses      54929    54880      -49     
- Partials    22776    22819      +43     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@adriangb
adriangb requested a review from kosiew October 10, 2026 18:10
@adriangb

Copy link
Copy Markdown
Contributor Author

@kosiew would you mind reviewing this change? Thanks!

Comment thread datafusion/common/src/cse.rs Outdated
// follow the same order.
node.apply_children(|child| {
self.conditional =
conditional.iter().any(|c| std::ptr::eq(*c, child));

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 kosiew left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@adriangb,

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`

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

done in b3004e5

@adriangb
adriangb enabled auto-merge October 12, 2026 12:55
@adriangb

Copy link
Copy Markdown
Contributor Author

Thanks for the review @kosiew !

@adriangb
adriangb added this pull request to the merge queue Oct 12, 2026
@github-merge-queue
github-merge-queue Bot removed this pull request from the merge queue due to no response for status checks Oct 12, 2026
@alamb
alamb added this pull request to the merge queue Oct 12, 2026
Merged via the queue into apache:main with commit e7d9971 Oct 12, 2026
42 checks passed
@adriangb
adriangb deleted the cse-aggregate-filter-args branch October 12, 2026 20:02
zzcclp pushed a commit to zzcclp/arrow-datafusion that referenced this pull request Oct 22, 2026
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]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

common Related to common crate optimizer Optimizer rules sqllogictest SQL Logic Tests (.slt) v56.0.0

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Common subexpression elimination evaluates the arguments of an aggregate with FILTER for rows that the filter rejects

5 participants