adriangb opened a new pull request, #26184:
URL: https://github.com/apache/datafusion/pull/26184
## Which issue does this PR close?
- Closes #26183.
## 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:
```sql
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 (#24507,
#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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
