vboo123 opened a new pull request, #59867:
URL: https://github.com/apache/spark/pull/59867
### What changes were proposed in this pull request?
Terms:
- The compile cache (`CodeGenerator.cache`) keeps each compiled generated
class by its source, so the same code is compiled once.
- A trial compile (`CodeCompiler.trial`) is how whole-stage codegen decides
whether to split a stage's expressions (SPARK-33301). It compiles the stage's
code in one piece, reads each method's bytecode size, and holds back what the
compile reports (the codegen metrics, the "Code generated in" log line and the
reports of methods past the JIT limit) until it knows which code it keeps.
Before this PR, the method sizes and the held reports were collected only
while a class was being compiled. A trial that found the class already in the
cache got neither, so `WholeStageCodegenExec.trialCompile` dropped the class
from the cache and compiled it again. That happens when the class is cached but
the split decision for its code is not: the decision was evicted, the same code
ran earlier with the split off or under another limit, or another thread
compiled it first.
This PR keeps the sizes and the reports with the class in the cache:
- Every compile into the cache runs under a trial of its own
(`CodeCompiler.held`). The cache keeps the class together with what that
compile held: the size of each method, and its reports in the order they were
made (`CodeCompiler.HeldCompile`).
- Inside a trial, `CodeGenerator.compile` gives the trial what the class's
compile held, whether it compiled the class or found it cached
(`TrialCompile.holdAll`). So a trial that keeps code it found cached still
reports it, as it did when it compiled the class again.
- `trialCompile` no longer drops and recompiles the class on a hit.
Outside a trial, a compile's reports are now made when the compile ends
instead of as they come, on the calling thread. With the JDK backend that is no
longer the compile worker thread, which has no task context or session conf.
### Why are the changes needed?
A trial that found its class cached compiled the stage's code a second time,
only to measure sizes the first compile had already measured. Dropping the
class could also make another thread's `doExecute`, which was about to use it,
compile it again. A thread that waited for another thread's compile of the same
class got no sizes at all, so the stage's decision was put off to a later run.
For example, in one session:
```scala
val branches = (0 until 300).map(i => s"WHEN id = $i THEN id *
$i").mkString(" ")
val q = s"SELECT CASE $branches ELSE -1 END AS r FROM range(100)"
spark.conf.set("spark.sql.codegen.wholeStage.splitExpressions", "false")
spark.sql(q).collect() // compiles the stage's code in one piece
spark.conf.set("spark.sql.codegen.wholeStage.splitExpressions", "true")
spark.sql(q).collect() // the trial finds that class in the cache
```
On master, the second `collect()` drops the cached class and compiles it
again before it decides to split. With this PR it decides from the sizes kept
in the cache.
### Does this PR introduce _any_ user-facing change?
No. Answers and generated code are the same, and the code a stage keeps
makes its metric updates and reports as before. Two small differences in what
is logged:
- When a trial keeps code it found in the cache, its "Code generated in"
line gives the time of that class's earlier compile.
- With the JDK backend on an executor, a method past the JIT limit is
reported at INFO, as `CodeCompiler.logHugeMethod` intends for executors.
Before, the report ran on the compile worker thread, which has no task context,
so it could be the once-per-JVM warning.
### How was this patch tested?
New test in `WholeStageCodegenSuite`, "SPARK-60049: a trial that finds the
class in the compile cache does not compile it". Under both Janino and the JDK
backend, it compiles a 300-branch stage's code outside a trial, then runs
`trialCompile` on the same code. It checks that the trial gets the method
sizes, that the cache still holds the same class instance (on master it holds a
new one, compiled again), and that the trial's held reports include the method
past the JIT limit.
The existing tests of the trial's reports in `WholeStageCodegenSuite` still
apply. "SPARK-33301: the code a stage keeps unsplit reports its method past the
JIT limit" goes through a trial that finds its class cached.
CI: GitHub Actions on my fork.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Kiro (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]
