MaxGekk opened a new pull request, #59225:
URL: https://github.com/apache/spark/pull/59225
### What changes were proposed in this pull request?
This replaces #59069. It is a redesign that takes in the reviews by
@dongjoon-hyun and @cloud-fan there.
Under whole-stage codegen,
`CodegenContext.splitExpressionsWithCurrentInputs` never split (`// TODO:
support whole stage codegen`), so every branch of a `CASE WHEN` went into the
operator's method. Outside a stage the branches are split into methods that
take the input row. Inside a stage the inputs are local variables of the
operator's method, and which of them a piece of code reads cannot be told from
the code strings. This PR splits a `CASE WHEN` inside a stage, and only when
the stage needs it.
**One input collector.** A private `CodegenContext.collectInputs` walks the
expressions a piece of code was generated from and returns the locals it reads,
or `None` when no method can take them. An `InputPolicy` records where its
three callers differ: the whole-stage split, `getLocalInputVariableValues`
(subexpressions, aggregates, `ExpandExec`), and `CommonExprSlots.methodArgs`
(`With`). It replaces the two walks that existed before and the third that
#59069 added. Under the whole-stage policy, the walk:
- takes the evaluated input variables of `currentVars`;
- takes `INPUT_ROW` only where the code reads it, including through a
`CommonExpressionRef` whose definition became a method taking the row (that
decision is now recorded when it is made);
- takes a subexpression state's value without descending below it;
- reads a slot of a compacted mutable state array as a field.
It refuses an input the operator has not evaluated yet, and a value no
parameter can carry. `getLocalInputVariableValues` keeps its signature and its
results.
**One split, with a per-block fallback.** `CaseWhen` passes each branch's
code together with the expressions it came from to a new `private[sql]
splitExpressionsWithSources`. The existing `splitExpressionsWithCurrentInputs`
keeps its signature for external callers.
- The code is packed into blocks by `methodSplitThreshold`, as everywhere,
and each block gets its own arguments.
- A block the collector refuses stays inline, in its place between the calls.
- Only when every block is refused does the `CASE WHEN` stay in one piece.
**One grouping.** The calls are grouped by the grouping of SPARK-59783
(#59042), now the only implementation, with a parameter for how many calls may
stay in the calling method. Inside a stage that is 8 calls. All of a
projection's outputs land in one method there, so a per-expression budget would
add up: four 300-branch `CASE WHEN`s gave a 13583-byte method that HotSpot
never compiles.
**The stage gate.** Splitting where the stage did not need it cost 5-15% at
64 branches on the runners: C2 compiles the unsplit method as one unit, but
stops inlining partway through the chain of calls. So a stage now decides from
its own code:
1. It generates its code without the split, recording which expressions it
could split.
2. It compiles that code.
3. If the compile fails, or a method is past
`spark.sql.codegen.wholeStage.splitExpressions.methodLimit` (8000 bytes by
default, HotSpot's limit for JIT-compiling a method), it generates the code
again with the split.
4. It keeps the split code only when that makes the stage's largest method
smaller and it compiles. Otherwise it runs the unsplit code.
A stage the JIT compiles whole therefore keeps exactly the code it had
before this PR. The details:
- The trial compile holds back its warnings and the huge-method report for
the code it discards, also on the JDK compiler's worker thread.
- A failed trial is remembered by a digest of the code, so re-executing a
query does not retry a 64 KB compile.
- The compile of the code that runs is cached as before.
**The discarded pass.** Subexpression elimination generates the common
subexpressions once without splitting, and again inside `subExpr_N` methods
when they are large. The functions added by the discarded first generation are
now removed from the class.
`spark.sql.codegen.wholeStage.splitExpressions` (internal, on by default)
turns all of this off. With it off, the generated code is exactly that of
master.
### Why are the changes needed?
SPARK-33301: a large `CASE WHEN` inside a stage fails to compile. On Spark
4.2.0, one `CASE WHEN` of `n` branches over a long column,
```scala
val branches = (1 to n).map(k => s"WHEN id = $k THEN id * $k").mkString(" ")
spark.sql(s"SELECT CASE $branches ELSE 0 END AS v FROM range(10)").collect()
```
compiles at 952 branches and fails at 953 ("Code grows beyond 64 KB"). Spark
then disables whole-stage codegen for the plan and runs it without. Much
earlier, from 96 branches in the benchmark's shape, the stage's method is past
8000 bytes, and HotSpot never JIT-compiles it (`-XX:+DontCompileHugeMethods`),
so the stage runs interpreted.
`CaseWhenCodegenBenchmark` (new), best times on GitHub runners, JDK 17 / 21
/ 25, at `42bbae0c829`. The results files of this PR's head are being generated
now and will be committed here.
| shape | split as needed | not split | whole stage codegen off |
|---|---|---|---|
| one CASE WHEN, 64 branches (6049 bytes) | 55 / 101 / 96 ms | 51 / 96 / 93
ms | 113 / 228 / 217 ms |
| one CASE WHEN, 96 branches (9057 bytes unsplit) | 70 / 136 / 126 ms | 1796
/ 3497 / 3966 ms | 152 / 284 / 298 ms |
| one CASE WHEN, 300 branches | 147 / 249 / 233 ms | 5752 / 10715 / 11973 ms
| 644 / 930 / 974 ms |
| a 300-branch CASE WHEN in a keyed aggregate | 181 / 308 / 276 ms | 5504 /
12218 / 12141 ms | 1605 / 2350 / 2239 ms |
Below 8000 bytes, "split as needed" runs the same code as "not split". The
difference there is within each rung's A/A spread ("not split" against "not
split again" in the files). Past it, the split is 25 to 51 times faster than
the unsplit stage and 2 to 9 times faster than running without whole-stage
codegen.
Four 300-branch `CASE WHEN`s in one projection ran 44-62% slower split as
needed than split always on those runs. The unsplit code fails to compile
there, and that commit retried the failed trial compile on every execution. The
digest memo in `484b4146d1b` removes the retry: on a laptop, a repeated plan
and compile of a 1000-branch stage went from 479 ms to 129 ms. The runs on the
head will show how much of the gap is left.
### Does this PR introduce _any_ user-facing change?
No, apart from the performance of such queries. Answers are the same. A
large `CASE WHEN` inside a stage now compiles, and keeps the stage's methods
under the JIT limit. The exceptions are blocks that read a column the operator
has not evaluated yet: they stay inline, and a `CASE WHEN` all of whose blocks
do so stays in one piece, as today. A stage the JIT already compiles whole is
unchanged.
### How was this patch tested?
New tests in `WholeStageCodegenSuite` (23) and `CollectInputsSuite` (10, one
per row of the collector's policies). Among them:
- **Splitting:** a 3000-branch `CASE WHEN` compiles inside a stage, and at
300 branches every method stays under 8000 bytes. Several large `CASE WHEN`s in
one projection keep its method small.
- **The gate:** a 64-branch stage keeps its unsplit code. The trial compile
reports nothing for the code it discards, and the kept code reports its huge
method. `chooseCode`'s decisions are checked with the compile faked, including
a split that fails to compile.
- **The per-block fallback:** a block reading an input not evaluated yet
stays inline between the calls. A block whose method would take too many
parameters stays inline too.
- **Inputs:** keyed aggregates, the buffer row and `max_by`. Sort-merge and
broadcast join conditions, with a common expression. The streamed-only
condition of an outer join, which a `CASE WHEN` never reaches. A `ROLLUP`, and
a split common subexpression. A slot of a compacted state array, read as a
field.
- **The discarded pass:** its functions are removed from the class.
Existing suites pass: `CodeGenerationSuite`, `ConditionalExpressionSuite`,
`SubexpressionEliminationSuite`, `WithExpressionEvalSuite`,
`WholeStageCodegenSuite`, the TPC-DS and TPC-H query suites, the join suites
and `DataFrameAggregateSuite`. Generated code for every TPC-DS query is
unchanged by the collector. The fork's full CI passed on `484b4146d1b`.
A differential fuzz, not part of the PR, compared random queries under
whole-stage codegen with the same queries without it. The table has nullable
numeric, string, date, decimal, array, struct and boolean columns, and the
queries cover ten operator shapes with nested `CASE WHEN`s in conditions and
values. It ran 72,000 queries over the method limit at 0, 1500 and the default,
both compilers, `splitConsumeFuncByOperator` on and off, and three seeds.
There was one difference, and it is not this PR's. With
`methodSplitThreshold=1`, a sort-merge join key holding a `coalesce` trips
SPARK-22668's testing assertion ("split function argument smj_streamedRow_0
cannot be a global variable"). This happens on master too:
`SortMergeJoinExec.createJoinKey` sets `INPUT_ROW` to the streamed row, which
is a field, and `Coalesce` splits through `splitExpressionsWithCurrentInputs`.
With this PR's split off, the same seed gives 32 differences, each a compile
failure past 64 KB.
Follow-ups, each for its own JIRA:
- the other callers of `splitExpressionsWithCurrentInputs` (`Coalesce`,
`In`, `Greatest`/`Least`, `CreateNamedStruct`, the hashes and others), passing
their sources;
- `Expression.reduceCodeSize`'s whole-stage TODO;
- letting `getLocalInputVariableValues` follow `With`;
- the join key assertion above.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Claude Opus 5.5)
--
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]
