xumingming opened a new pull request, #57703:
URL: https://github.com/apache/spark/pull/57703
### What changes were proposed in this pull request?
Wire `ExpandExec` into the whole-stage codegen subexpression elimination
framework:
- All branch expressions are bound once up front and analyzed together. A
subexpression shared across branches (or repeated within a branch) is evaluated
once per input row, before the branch loop, and each branch references the
cached value. This is semantics-preserving because every branch of an Expand
consumes the same input row.
- Both the branch-invariant column generation and the per-branch switch/case
generation resolve repeated subtrees to the cached values.
- The switch/case function splitting (SPARK-35329) now passes the eliminated
subexpression variables into the split functions as parameters, so large
Expands keep compiling correctly.
- Gated by the existing `spark.sql.subexpressionElimination.enabled` conf
(no new configuration).
- Adds an `ExpandBenchmark` case.
### Why are the changes needed?
Conditional-aggregate rollup queries (e.g. "N-day active users" dashboards)
stack many conditional aggregates whose conditions share one expensive
subexpression:
```sql
SELECT
COUNT(DISTINCT IF(datediff(date '2026-01-01',
from_unixtime(unix_timestamp(ts, 'yyyy-MM-dd HH:mm:ss'))) <= 1, uid,
NULL)) AS uv_1d,
COUNT(DISTINCT IF(datediff(date '2026-01-01',
from_unixtime(unix_timestamp(ts, 'yyyy-MM-dd HH:mm:ss'))) <= 7, uid,
NULL)) AS uv_7d,
SUM(IF(datediff(date '2026-01-01',
from_unixtime(unix_timestamp(ts, 'yyyy-MM-dd HH:mm:ss'))) <= 1, 1, 0))
AS pv_1d
-- ... more conditional aggregates over longer windows
FROM traffic
```
`RewriteDistinctAggregates` gives each distinct group its own Expand branch
with the condition expression verbatim, so the shared subexpression is compiled
into every branch body and re-evaluated once per branch per input row. In the
benchmark query above it is evaluated 18 times per input row (9
distinct-aggregate conditions + 9 regular-aggregate conditions across 10
branches); with this PR it is evaluated once.
### Does this PR introduce _any_ user-facing change?
No. Query results are unchanged; only the number of evaluations per input
row changes.
### How was this patch tested?
- Extended the `WholeStageCodegenSuite` test "Expand should eliminate common
subexpressions across branches" (SQL-based, via conditional `COUNT DISTINCT`
aggregates):
- asserts the shared subexpression is evaluated once per input row with
elimination enabled, and once per branch with it disabled;
- forces both switch/case code paths deterministically via
`spark.sql.codegen.methodSplitThreshold` (a tiny threshold forces the
SPARK-35329 function splitting, a huge threshold keeps the bodies inline),
asserting on the generated `switchCaseCode` functions and verifying correctness
on both paths.
- Ran the full `WholeStageCodegenSuite` (57 tests passed).
- Benchmark: new `ExpandBenchmark` case modeling the rollup above, 5M rows,
3 iterations, Apple M4 Pro / JDK 17 (local run): shared subexpression
evaluations per input row 18 -> 1; average runtime 82485 ms -> 51457 ms (15701
-> 9756 ns/row), 1.6X faster.
### Was this patch authored or co-authored using generative AI tooling?
No.
--
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]
