MaxGekk opened a new pull request, #59069:
URL: https://github.com/apache/spark/pull/59069
### What changes were proposed in this pull request?
Under whole-stage codegen,
`CodegenContext.splitExpressionsWithCurrentInputs` did not split at all (`//
TODO: support whole stage codegen`), so the code of every branch of a `CASE
WHEN` went into the operator's method. Outside a stage the same code is 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 is
not known from the code strings the method receives. This PR teaches it to
split under whole-stage codegen when the caller also passes the expressions the
code was generated for, and makes `CaseWhen` pass them.
From those expressions, a new private method,
`CodegenContext.wholeStageSplitArguments`, computes the arguments of the split
methods: the local variables the code reads from the operator's method. It
walks the expressions the way the existing split mechanisms under whole-stage
codegen do. It takes the evaluated input variables of `currentVars` and
`INPUT_ROW` where an operator sets one, as
`CodeGenerator.getLocalInputVariableValues` does for aggregates, subexpressions
and `ExpandExec`. It takes the variables of a subexpression elimination state
without descending below it, as that method does too, so that the split
methods' arguments are among those of an enclosing method an aggregate or
`ExpandExec` has split out, where the calls then sit. And it follows a
`CommonExpressionRef` into its definition, as `CommonExprSlots.methodArgs` does
for `With`. Fields and literals need no argument.
It refuses, and the code stays in one piece as before, where a local the
code reads cannot be named or passed:
- an input the operator has not evaluated yet (a projection defers reading a
column it uses once): its code is inside this code and reads locals of the
operator producing the row, as `methodArgs` explains;
- a value no parameter can carry: a `SimpleExprValue`, or a variable naming
a slot of a compacted state array;
- a subexpression state for `Alias`, `Collate` or an identity `Cast`, whose
`genCode` generates the child instead of reading the state;
- a `LambdaVariable`, whose value is a local of an object expression's loop;
- more parameter slots than a method can take.
Callers that do not pass the expressions keep today's behavior. The split is
on by default and can be turned off with the internal
`spark.sql.codegen.wholeStage.splitExpressions`, as
`spark.sql.codegen.aggregate.splitAggregateFunc.enabled` can for the aggregate
split.
### Why are the changes needed?
SPARK-33301: a large `CASE WHEN` fails to compile inside a stage. 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
logs "Failed to compile the generated Java code" and "Whole-stage codegen
disabled for plan" and runs the stage without whole-stage codegen. Earlier,
from 138 branches, the stage's method is past 8000 bytes and HotSpot never
JIT-compiles it (`-XX:+DontCompileHugeMethods`), so the stage runs interpreted.
With this change the branches are split into methods at both sizes.
### Does this PR introduce _any_ user-facing change?
No, other than the performance of such queries. The answers are the same; a
large `CASE WHEN` inside a stage now compiles and is JIT-compiled.
### How was this patch tested?
New tests in `WholeStageCodegenSuite`:
- a 3000-branch `CASE WHEN` compiles inside a stage and answers (before: the
compile error, thrown under testing);
- at 300 branches the stage's methods stay under 8000 bytes, and with the
config off the method is past 8000 bytes;
- the split methods take nullable and string columns and subexpression
locals (`upper(s)` in every condition);
- they take the aggregation buffer row of a keyed aggregate, called from
inside the aggregate's own split method;
- a `CASE WHEN` that is itself a split common subexpression;
- a `CASE WHEN` reading a column the projection has not evaluated stays in
one piece.
Existing suites: `CodeGenerationSuite`, `ConditionalExpressionSuite`,
`SubexpressionEliminationSuite`, `CodegenSubexpressionEliminationSuite` and
`WithExpressionEvalSuite` in catalyst; `WholeStageCodegenSuite`,
`TPCDSQuerySuite`, `TPCDSQueryWithStatsSuite`, `TPCDSQueryANSISuite`,
`TPCHQuerySuite` and `DataFrameAggregateSuite` in sql. No TPC-DS stage
generates code differently: none has a multi-branch `CASE WHEN` inside a stage
(a single-branch one is generated as `If`).
Beyond the unit tests, with a temporary change (not in this PR) that splits
every `CASE WHEN` inside a stage into one method per branch, so that the
argument computation is exercised on every such expression:
- `SQLQueryTestSuite`, `SQLQuerySuite`, the DataFrame and Dataset suites,
`DataFrameAggregateSuite`, `JoinSuite` and `WholeStageCodegenSuite` pass (3162
tests; 240 `CASE WHEN`s split, 11 refused), apart from `udtf/udtf.sql`, whose
Python error message differs by Python version here;
- a differential fuzz of 600 random queries over a Parquet table (nullable
numeric, string, date, decimal, array, struct and boolean columns) and nine
operator shapes -- a projection, a repeated `CASE WHEN` (a split common
subexpression), keyed and unkeyed aggregates, a filter, broadcast and
sort-merge joins, `explode`, `ROLLUP`, a `CASE WHEN` inside a lambda, `ORDER BY
... LIMIT` -- with nested `CASE WHEN`s in conditions and values, compared
against the same queries without whole-stage codegen: no difference, no compile
error, 2343 splits and 430 refusals (all of them an input the operator had not
evaluated). 300 more at the default threshold: no difference.
This PR builds on SPARK-59783 (#59042, merged), which groups the calls to
split functions: without it, at 1000 branches the method holding the calls to
the split methods is itself past 8000 bytes and runs interpreted. With it, on a
laptop, the 1000-branch query takes about 0.36 s a run in a split stage against
1.73 s without whole-stage codegen; at 300 branches the split stage is 2.8x
faster than without whole-stage codegen and 22x faster than today's unsplit
stage, whose method HotSpot never compiles. A split `CASE WHEN` inside a common
subexpression that is split is generated twice by subexpression elimination, so
the first generation's methods are unused, as for any code subexpression
elimination regenerates.
### 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]
