Repository navigation
Commit a1b272a
committed
[SPARK-33301][SQL] Split a large CASE WHEN into methods under whole-stage codegen
### What changes were proposed in this pull request?
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. With the split on this holds under every policy and in every stage: a slot counts as the field it is when code moves into a method. An operator's method that took it as a parameter, and failed to compile, now compiles; a `With` definition reading one, which stayed inline, gets its `computeCommonExpr` method. This fixes a stage whose large common subexpression reads such a slot, such as an UNPIVOT's string output: on master its `subExpr` method takes `UTF8String expand_mutableStateArray_0[1]` as a parameter, fails to compile, and the stage falls back;
- descends below a node whose code is its child's (`Alias`, `Collate`, an identity `Cast`) instead of refusing its state.
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 decides per method, from its own code:
1. It generates its code in one piece, marking the code of each expression a split could take.
2. It compiles that code as a trial, which reports every method's bytecode size.
3. The methods run per row that are past `spark.sql.codegen.wholeStage.splitExpressions.methodLimit` (8000 bytes by default, HotSpot's limit for JIT-compiling a method) or `spark.sql.codegen.hugeMethodLimit` are the ones a split can help. It generates the code again with only the expressions marked in those methods split. An expression in a method the JIT compiles whole is left alone.
4. It keeps the split code when that lowers the bytes past the limit and it compiles; otherwise it runs the code in one piece. If the code in one piece does not compile at all, every expression is split where that compiles; where it does not either, the code in one piece goes to `doExecute`, which reports, falls back or throws exactly as on master.
A stage the JIT compiles whole therefore keeps the code it had before this PR, including the helper functions of a discarded subexpression pass, except where code moving into a method reads a compacted slot: it reads it as a field (see above). That workaround can go once SPARK-60021 labels `ExpandExec`'s outputs as the fields they are. The details:
- The decision is remembered by a digest of the marked code, with the class loader, the compiler, the limit and `methodSplitThreshold`, so re-executing a query skips the trial compiles. Threads that meet the stage at once wait for one decision, under a per-key lock. A split that fails to compile is warned about, its code kept in one piece, thrown under testing, and remembered; the warning carries the error only where the split's compile did not log it already; where a trial's method sizes are unknown (a compile cache hit), nothing is.
- The split calls are matched between the two generations by their expressions, with lambda variables' values ignored, and their occurrence, so a copy of a CASE WHEN in a method under the limit is not split with the one past it.
- A split generation that turns no block into a method is not compiled; the code in one piece is kept, as where no split is tried: its trial's class stays cached for `doExecute` where it compiled, and where it did not, `doExecute` compiles it again and reports the failure as on master.
- The variant not kept leaves the compile cache, so the cache holds only code that runs, and a later compile of that variant reports what it finds.
- The trial compile holds back the huge-method reports, the "Code generated in" line and the codegen metrics for the code it discards, also on the JDK compiler's worker thread, and logs an expected failure at DEBUG. The `spark.sql.codegen.logLevel` source dump is not held back, so with it on, each variant compiled is dumped.
- Recording does not pack blocks, so below the limit the gate costs about 0.04 ms per code generation on a 64-branch stage (1.60 ms against 1.56 ms without the split).
**The discarded pass.** Subexpression elimination generates the common subexpressions once without splitting, and again inside `subExpr_N` methods when they are large. A generation that splits the stage's expressions now removes the functions added by the discarded first generation. A stage kept in one piece keeps them, as on master.
`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; its three results files are committed), best times on GitHub runners. JDK 17 ran on an AMD EPYC 7763, JDK 21 on an AMD EPYC 9V74 and JDK 25 on an Intel Xeon Platinum 8573C, so the columns compare the arms within a JDK, not the JDKs with each other.
| shape | split as needed | not split | split always | whole stage codegen off |
|---|---|---|---|---|
| one CASE WHEN, 64 branches (6049 bytes) | 92 / 95 / 63 ms | 92 / 94 / 64 ms | 97 / 100 / 64 ms | 211 / 219 / 152 ms |
| one CASE WHEN, 96 branches (9057 bytes unsplit) | 120 / 122 / 82 ms | 3992 / 3717 / 2394 ms | 115 / 115 / 75 ms | 296 / 305 / 200 ms |
| one CASE WHEN, 300 branches | 249 / 227 / 162 ms | 11854 / 12083 / 7043 ms | 232 / 216 / 143 ms | 976 / 951 / 709 ms |
| four 300-branch CASE WHENs in one projection | 1099 / 868 / 589 ms | 4638 / 5173 / 3457 ms | 1042 / 760 / 505 ms | 4643 / 5172 / 3451 ms |
| a 300-branch CASE WHEN in a keyed aggregate | 262 / 270 / 191 ms | 12558 / 11943 / 6987 ms | 243 / 234 / 167 ms | 2156 / 2358 / 2165 ms |
Below 8000 bytes, "split as needed" runs the same code as "not split", and times the same: each arm's repeat ("again" in the files) is within 1 ms of it. Past it, the split is 30 to 50 times faster than the unsplit stage, which HotSpot never compiles, and 2 to 11 times faster than running without whole-stage codegen.
What splitting as needed still costs against "split always" is the code generation it repeats: a stage that splits generates its code twice per execution, once in one piece and once split. That is 5-13% on one 300-branch CASE WHEN, 5-17% on four of them in one projection, and 8-15% in the keyed aggregate. The decision is remembered by the generated code, so a query whose code carries a literal that changes on every run, such as the timestamp `current_timestamp()` becomes, decides afresh each run and pays two generations and two compiles where master pays one of each; queries without such literals compile nothing new after the first execution. Remembering the decision by a key that ignores literal values would remove both; it is left for a follow-up (SPARK-60048).
### 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, except where code moving into a method reads a compacted slot.
### How was this patch tested?
New tests in `WholeStageCodegenSuite` (35), `CollectInputsSuite` (12, one per row of the collector's policies, and the conf off) and `CodeCompilerSuite` (1, the scanner of generated Java). 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, byte for byte, also with common subexpressions. A stage splits only the CASE WHENs in its methods past the limit: a filter's 64-branch CASE WHEN in `processNext` stays whole beside a projection's 300-branch one that splits. The trial compile reports no huge method, logs no compile time and updates no metric for the code it discards, which each test of `chooseCode`'s decisions pins by which trial's held reports it replays; the kept code reports its huge method, also after the same code was discarded in an earlier trial. `chooseCode`'s decisions are checked with the compile faked: a split kept, a larger unsplittable method beside, a split that turns no block into a method, a split that fails to compile, warned about or thrown, code in one piece that fails to compile with the split compiling, failing or splitting nothing, unknown sizes not remembered, the memo, and two threads meeting one undecided key, which decide once. With real generations, a split that turns no block into a method, beside a discarded subexpression pass's helpers, compiles nothing more. A query run at two `methodSplitThreshold`s in one session, in both orders, splitting the right CASE WHEN. Only the kept code counts in the codegen metrics.
- **Markers:** a literal holding a marker's text is left as it is, also in the code kept in one piece.
- **Matching:** a pushed-down filter's copy of a projection's CASE WHEN stays whole, and a CASE WHEN in a join key whose branches hold a `transform`'s lambda variables is split, its method being past the JIT limit unsplit.
- **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, and a CASE WHEN inside a CASE WHEN's branch whose common expression takes the build side's row, where the outer split must pass the row on. 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 under every input policy with the split on; with it off, an operator's method takes it as a parameter, and a `With` definition reading one stays inline. A simple CASE over a collated column, and over a cast that only renames struct fields. An UNPIVOT whose large common subexpression reads a compacted slot: with the conf on, its `subExpr` method reads the slot as a field and compiles; with it off, the method takes the slot as a parameter, as on master, checked on the generated code so that the test holds under either compiler backend.
- **The discarded pass:** a generation that splits removes its functions from the class; one kept in one piece keeps them.
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.
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, then 24,000 more on the per-method gate. The second run found a split that did not compile, the nested CASE WHEN above, which is fixed and pinned by a test; after the fix the runs agree on every query but one.
That one difference 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 with its own JIRA:
- SPARK-60021: label `ExpandExec`'s outputs as the fields they are. The compacted-slot rule of this PR is then redundant, and goes in the same change if it is on master by then;
- SPARK-60046: the join key assertion above;
- SPARK-60047: the other callers of `splitExpressionsWithCurrentInputs` (`Coalesce`, `In`, `Greatest`/`Least`, `CreateNamedStruct`, the hashes and others) passing their sources, and `Expression.reduceCodeSize`'s whole-stage TODO;
- SPARK-60048: a decision key that ignores literal values;
- SPARK-60049: keeping each method's size with `ByteCodeStats` in the compile cache, so that a cache hit needs no second trial compile.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Claude Opus 5.5)
Closes #59225 from MaxGekk/SPARK-33301-redesign.
Authored-by: Max Gekk <max.gekk@gmail.com>
Signed-off-by: Max Gekk <max.gekk@gmail.com>
(cherry picked from commit 912c707)
Signed-off-by: Max Gekk <max.gekk@gmail.com>1 parent 320f236 commit a1b272a
15 files changed
Lines changed: 3591 additions & 250 deletions
File tree
- sql
- catalyst/src
- main/scala/org/apache/spark/sql
- catalyst/expressions
- codegen
- internal
- test/scala/org/apache/spark/sql/catalyst/expressions/codegen
- core
- benchmarks
- src
- main/scala/org/apache/spark/sql/execution
- aggregate
- test/scala/org/apache/spark/sql/execution
- benchmark
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
1618 | 1618 | | |
1619 | 1619 | | |
1620 | 1620 | | |
1621 | | - | |
| 1621 | + | |
1622 | 1622 | | |
1623 | 1623 | | |
1624 | 1624 | | |
1625 | 1625 | | |
1626 | 1626 | | |
1627 | 1627 | | |
| 1628 | + | |
| 1629 | + | |
| 1630 | + | |
| 1631 | + | |
1628 | 1632 | | |
1629 | 1633 | | |
1630 | 1634 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
26 | 26 | | |
27 | 27 | | |
28 | 28 | | |
| 29 | + | |
29 | 30 | | |
30 | 31 | | |
31 | 32 | | |
| |||
326 | 327 | | |
327 | 328 | | |
328 | 329 | | |
329 | | - | |
330 | | - | |
| 330 | + | |
| 331 | + | |
| 332 | + | |
| 333 | + | |
| 334 | + | |
331 | 335 | | |
332 | 336 | | |
333 | 337 | | |
334 | 338 | | |
335 | 339 | | |
336 | 340 | | |
| 341 | + | |
| 342 | + | |
| 343 | + | |
| 344 | + | |
| 345 | + | |
| 346 | + | |
| 347 | + | |
| 348 | + | |
| 349 | + | |
| 350 | + | |
| 351 | + | |
| 352 | + | |
| 353 | + | |
| 354 | + | |
| 355 | + | |
| 356 | + | |
| 357 | + | |
| 358 | + | |
| 359 | + | |
| 360 | + | |
| 361 | + | |
| 362 | + | |
| 363 | + | |
| 364 | + | |
| 365 | + | |
| 366 | + | |
| 367 | + | |
| 368 | + | |
| 369 | + | |
| 370 | + | |
| 371 | + | |
| 372 | + | |
| 373 | + | |
| 374 | + | |
| 375 | + | |
| 376 | + | |
| 377 | + | |
| 378 | + | |
| 379 | + | |
| 380 | + | |
| 381 | + | |
| 382 | + | |
| 383 | + | |
| 384 | + | |
| 385 | + | |
| 386 | + | |
| 387 | + | |
| 388 | + | |
| 389 | + | |
| 390 | + | |
| 391 | + | |
| 392 | + | |
| 393 | + | |
| 394 | + | |
| 395 | + | |
| 396 | + | |
| 397 | + | |
| 398 | + | |
| 399 | + | |
| 400 | + | |
| 401 | + | |
| 402 | + | |
| 403 | + | |
| 404 | + | |
| 405 | + | |
| 406 | + | |
| 407 | + | |
| 408 | + | |
| 409 | + | |
| 410 | + | |
| 411 | + | |
| 412 | + | |
| 413 | + | |
| 414 | + | |
| 415 | + | |
| 416 | + | |
| 417 | + | |
| 418 | + | |
| 419 | + | |
| 420 | + | |
| 421 | + | |
| 422 | + | |
| 423 | + | |
| 424 | + | |
| 425 | + | |
| 426 | + | |
| 427 | + | |
337 | 428 | | |
338 | 429 | | |
339 | 430 | | |
| |||
342 | 433 | | |
343 | 434 | | |
344 | 435 | | |
345 | | - | |
| 436 | + | |
| 437 | + | |
346 | 438 | | |
347 | 439 | | |
348 | 440 | | |
| |||
378 | 470 | | |
379 | 471 | | |
380 | 472 | | |
381 | | - | |
| 473 | + | |
382 | 474 | | |
383 | 475 | | |
384 | 476 | | |
385 | 477 | | |
386 | 478 | | |
387 | 479 | | |
388 | | - | |
| 480 | + | |
| 481 | + | |
389 | 482 | | |
390 | | - | |
| 483 | + | |
| 484 | + | |
| 485 | + | |
| 486 | + | |
| 487 | + | |
391 | 488 | | |
392 | 489 | | |
393 | 490 | | |
| |||
423 | 520 | | |
424 | 521 | | |
425 | 522 | | |
| 523 | + | |
| 524 | + | |
| 525 | + | |
426 | 526 | | |
427 | 527 | | |
428 | 528 | | |
429 | 529 | | |
430 | 530 | | |
431 | 531 | | |
432 | 532 | | |
| 533 | + | |
| 534 | + | |
| 535 | + | |
| 536 | + | |
| 537 | + | |
| 538 | + | |
| 539 | + | |
| 540 | + | |
| 541 | + | |
| 542 | + | |
| 543 | + | |
| 544 | + | |
| 545 | + | |
| 546 | + | |
| 547 | + | |
| 548 | + | |
| 549 | + | |
| 550 | + | |
| 551 | + | |
| 552 | + | |
| 553 | + | |
| 554 | + | |
| 555 | + | |
| 556 | + | |
| 557 | + | |
| 558 | + | |
| 559 | + | |
| 560 | + | |
| 561 | + | |
| 562 | + | |
| 563 | + | |
| 564 | + | |
| 565 | + | |
| 566 | + | |
| 567 | + | |
| 568 | + | |
| 569 | + | |
| 570 | + | |
| 571 | + | |
| 572 | + | |
| 573 | + | |
| 574 | + | |
| 575 | + | |
| 576 | + | |
| 577 | + | |
| 578 | + | |
| 579 | + | |
| 580 | + | |
| 581 | + | |
| 582 | + | |
| 583 | + | |
| 584 | + | |
| 585 | + | |
| 586 | + | |
| 587 | + | |
| 588 | + | |
| 589 | + | |
| 590 | + | |
| 591 | + | |
| 592 | + | |
433 | 593 | | |
434 | 594 | | |
435 | 595 | | |
| |||
472 | 632 | | |
473 | 633 | | |
474 | 634 | | |
475 | | - | |
| 635 | + | |
476 | 636 | | |
477 | 637 | | |
478 | 638 | | |
479 | | - | |
| 639 | + | |
480 | 640 | | |
481 | 641 | | |
482 | 642 | | |
| |||
697 | 857 | | |
698 | 858 | | |
699 | 859 | | |
700 | | - | |
| 860 | + | |
701 | 861 | | |
702 | 862 | | |
703 | 863 | | |
704 | | - | |
705 | | - | |
706 | 864 | | |
707 | 865 | | |
708 | 866 | | |
709 | | - | |
710 | | - | |
711 | | - | |
712 | | - | |
713 | | - | |
714 | | - | |
715 | | - | |
716 | | - | |
717 | | - | |
718 | | - | |
719 | | - | |
720 | | - | |
721 | | - | |
722 | | - | |
723 | | - | |
724 | | - | |
725 | | - | |
726 | | - | |
727 | | - | |
728 | | - | |
729 | | - | |
730 | | - | |
731 | | - | |
732 | | - | |
733 | | - | |
734 | | - | |
735 | | - | |
736 | | - | |
737 | | - | |
738 | | - | |
739 | | - | |
740 | | - | |
741 | | - | |
742 | | - | |
743 | | - | |
744 | | - | |
745 | | - | |
746 | | - | |
747 | | - | |
748 | | - | |
749 | | - | |
750 | | - | |
751 | | - | |
752 | | - | |
753 | | - | |
754 | | - | |
755 | | - | |
756 | | - | |
757 | | - | |
758 | | - | |
759 | | - | |
| 867 | + | |
| 868 | + | |
| 869 | + | |
| 870 | + | |
| 871 | + | |
760 | 872 | | |
761 | | - | |
762 | 873 | | |
763 | 874 | | |
764 | 875 | | |
| |||
1316 | 1427 | | |
1317 | 1428 | | |
1318 | 1429 | | |
| 1430 | + | |
| 1431 | + | |
1319 | 1432 | | |
1320 | | - | |
| 1433 | + | |
1321 | 1434 | | |
| 1435 | + | |
1322 | 1436 | | |
1323 | 1437 | | |
1324 | 1438 | | |
| |||
1370 | 1484 | | |
1371 | 1485 | | |
1372 | 1486 | | |
1373 | | - | |
| 1487 | + | |
1374 | 1488 | | |
1375 | 1489 | | |
1376 | 1490 | | |
| |||
1389 | 1503 | | |
1390 | 1504 | | |
1391 | 1505 | | |
1392 | | - | |
| 1506 | + | |
1393 | 1507 | | |
1394 | 1508 | | |
1395 | 1509 | | |
| |||
0 commit comments