Visitar URL original
[SPARK-33301][SQL] Split a large CASE WHEN into methods under whole-s… · apache/spark@a1b272a · GitHub
Skip to content

Commit a1b272a

Browse files
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/Cast.scala‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1618,13 +1618,17 @@ case class Cast(
16181618
override def genCode(ctx: CodegenContext): ExprCode = {
16191619
// If the cast does not change the structure, then we don't really need to cast anything.
16201620
// We can return what the children return. Same thing should happen in the interpreted path.
1621-
if (DataType.equalsStructurally(child.dataType, dataType)) {
1621+
if (generatesChildCode) {
16221622
child.genCode(ctx)
16231623
} else {
16241624
super.genCode(ctx)
16251625
}
16261626
}
16271627

1628+
/** Whether `genCode` generates the child's code as this cast's, the cast changing nothing. */
1629+
private[catalyst] def generatesChildCode: Boolean =
1630+
DataType.equalsStructurally(child.dataType, dataType)
1631+
16281632
override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
16291633
val eval = child.genCode(ctx)
16301634
val nullSafeCast = nullSafeCastFunction(child.dataType, dataType, ctx)

‎sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeCompiler.scala‎

Lines changed: 180 additions & 66 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ import javax.tools.{Diagnostic, DiagnosticCollector, FileObject, ForwardingJavaF
2626

2727
import scala.collection.mutable
2828
import scala.jdk.CollectionConverters._
29+
import scala.util.Try
2930
import scala.util.control.NonFatal
3031

3132
import com.google.common.cache.{Cache, CacheBuilder}
@@ -326,14 +327,104 @@ object CodeCompiler extends Logging {
326327
private lazy val jitSkipsHugeMethods: Boolean =
327328
Utils.getVMOptionValue("DontCompileHugeMethods").contains("true")
328329

329-
/** Generated methods that run once per class or per partition, not per row. */
330-
private val runOnceMethods = Set("<init>", "<clinit>", "init")
330+
/**
331+
* Generated methods that run once per class or per partition, not per row, so that one being
332+
* interpreted costs nothing per row.
333+
*/
334+
private[sql] val runOnceMethods = Set("<init>", "<clinit>", "init")
331335

332336
private val hugeMethodWarned = new java.util.concurrent.atomic.AtomicBoolean(false)
333337

334338
/** Lets a test see the first warning again. */
335339
private[catalyst] def resetHugeMethodWarning(): Unit = hugeMethodWarned.set(false)
336340

341+
/**
342+
* A compile whose result its caller may discard ([[trial]]). What the compile would report of
343+
* the code - its updates of the codegen metrics, which describe the code that runs, its "Code
344+
* generated in" log line with the compile time, and its reports of methods past the JIT limit -
345+
* is held here, in the order made, for the caller to make if it keeps the code.
346+
*
347+
* @param failureExpected whether a failure to compile is an answer the caller asks for, which the
348+
* compile then logs at debug level rather than as an error.
349+
*/
350+
private[sql] final class TrialCompile private[codegen] (val failureExpected: Boolean) {
351+
private val methods = mutable.ArrayBuffer.empty[(String, Int)]
352+
private val reports = mutable.ArrayBuffer.empty[() => Unit]
353+
354+
private[codegen] def hold(report: () => Unit): Unit = reports += report
355+
356+
private[codegen] def recordMethod(methodName: String, size: Int): Unit =
357+
methods += ((methodName, size))
358+
359+
/**
360+
* The bytecode size of every method the compile produced, by its name in the source, the
361+
* largest where a name is in more than one class; empty where the compile was a hit in the
362+
* compile cache, which produces nothing. Janino names a private method of an inner class with
363+
* a `$` after its name in the source.
364+
*/
365+
def methodSizes: Map[String, Int] =
366+
methods.groupMapReduce(_._1.stripSuffix("$"))(_._2)(math.max)
367+
368+
/** Makes the reports and metric updates the compile held back, for the code kept. */
369+
def report(): Unit = reports.foreach(_())
370+
}
371+
372+
private val currentTrial = new ThreadLocal[TrialCompile]
373+
374+
/**
375+
* Runs `body`, a compile whose result its caller may discard, holding back its codegen metric
376+
* updates, its compile time log line and its reports of methods past the JIT limit until the
377+
* caller keeps the code ([[TrialCompile.report]]); and, where `failureExpected`, logging its
378+
* failure to compile at debug level, the failure being then the answer the caller is asking
379+
* for. The source dump of `spark.sql.codegen.logLevel` is made for every compile, a trial's
380+
* included. Whole-stage codegen compiles a stage this way to decide whether to split its
381+
* expressions (`spark.sql.codegen.wholeStage.splitExpressions.methodLimit`). The trial is the
382+
* calling thread's, and a backend that compiles on a thread of its own carries it there
383+
* ([[withTrial]]), as the JDK backend does.
384+
*/
385+
private[sql] def trial[T](failureExpected: Boolean)(body: => T): (Try[T], TrialCompile) = {
386+
val trial = new TrialCompile(failureExpected)
387+
(withTrial(trial)(Try(body)), trial)
388+
}
389+
390+
/** The trial the calling thread compiles under, or null outside one. */
391+
private[codegen] def activeTrial: TrialCompile = currentTrial.get
392+
393+
/**
394+
* Makes `report`, a metric update or a log line about the code compiled - its compile time, or
395+
* a method past the JIT limit - or, in a trial, holds it back until the caller keeps the code
396+
* ([[TrialCompile.report]]), so that code compiled only to be discarded is neither counted in
397+
* the metrics nor reported.
398+
*/
399+
private[codegen] def reportOrHold(report: => Unit): Unit = Option(currentTrial.get) match {
400+
case Some(trial) => trial.hold(() => report)
401+
case None => report
402+
}
403+
404+
/** Runs `body` on this thread under `trial`, which is null outside one. */
405+
private[codegen] def withTrial[T](trial: TrialCompile)(body: => T): T = {
406+
val previous = currentTrial.get
407+
currentTrial.set(trial)
408+
try body finally currentTrial.set(previous)
409+
}
410+
411+
private def failureExpected: Boolean = {
412+
val trial = currentTrial.get
413+
trial != null && trial.failureExpected
414+
}
415+
416+
/**
417+
* Reports a failure to compile generated code, at debug level in a trial that expects one
418+
* ([[trial]]).
419+
*/
420+
private[codegen] def logCompileFailure(e: Throwable): Unit = {
421+
if (failureExpected) {
422+
logDebug("Failed to compile the generated Java code of a trial compile.", e)
423+
} else {
424+
logError("Failed to compile the generated Java code.", e)
425+
}
426+
}
427+
337428
/**
338429
* Reports a generated method whose bytecode is past HotSpot's JIT limit
339430
* ([[CodeGenerator.DEFAULT_JVM_HUGE_METHOD_LIMIT]]). The first time in the driver that such a
@@ -342,7 +433,8 @@ object CodeCompiler extends Logging {
342433
* once, and `spark.sql.codegen.hugeMethodLimit` does not already make the stage fall back.
343434
* Every other report is INFO: every one after that first warning, every one from a generated
344435
* class outside whole-stage codegen, where the setting does not apply, and every one on an
345-
* executor, since the driver compiles each stage before its tasks do and has reported it.
436+
* executor, since the driver compiles each stage before its tasks do and has reported it. In a
437+
* trial, its caller holds the report back with the metric updates ([[reportOrHold]]).
346438
*/
347439
private[catalyst] def logHugeMethod(className: String, methodName: String, size: Int): Unit = {
348440
val limit = CodeGenerator.DEFAULT_JVM_HUGE_METHOD_LIMIT
@@ -378,16 +470,21 @@ object CodeCompiler extends Logging {
378470
classBytecodes: Iterable[(String, Array[Byte])]): ByteCodeStats = {
379471
val perClass = classBytecodes.map { case (_, classBytes) =>
380472
val classCodeSize = classBytes.length
381-
CodegenMetrics.METRIC_GENERATED_CLASS_BYTECODE_SIZE.update(classCodeSize)
473+
reportOrHold(CodegenMetrics.METRIC_GENERATED_CLASS_BYTECODE_SIZE.update(classCodeSize))
382474
try {
383475
val cf = new ClassFile(new ByteArrayInputStream(classBytes))
384476
val constPoolSize = cf.getConstantPoolSize
385477
val methodCodeSizes = cf.methodInfos.asScala.flatMap { method =>
386478
method.getAttributes.collect { case attr: CodeAttribute =>
387479
val byteCodeSize = attr.code.length
388-
CodegenMetrics.METRIC_GENERATED_METHOD_BYTECODE_SIZE.update(byteCodeSize)
480+
reportOrHold(CodegenMetrics.METRIC_GENERATED_METHOD_BYTECODE_SIZE.update(byteCodeSize))
481+
Option(currentTrial.get).foreach(_.recordMethod(method.getName, byteCodeSize))
389482
if (byteCodeSize > CodeGenerator.DEFAULT_JVM_HUGE_METHOD_LIMIT) {
390-
logHugeMethod(cf.getThisClassName, method.getName, byteCodeSize)
483+
// Bound first, so that a report held in a trial keeps the names, not the parsed
484+
// class.
485+
val className = cf.getThisClassName
486+
val methodName = method.getName
487+
reportOrHold(logHugeMethod(className, methodName, byteCodeSize))
391488
}
392489
byteCodeSize
393490
}
@@ -423,13 +520,76 @@ object CodeCompiler extends Logging {
423520
* the default conf instead of the calling session's.
424521
*/
425522
private[codegen] def logGeneratedCodeOnFailure(code: CodeAndComment, maxLines: Int): Unit = {
523+
if (failureExpected) {
524+
return
525+
}
426526
val formatted = s"\n${CodeFormatter.format(code, maxLines)}"
427527
if (Utils.isTesting) {
428528
logError(formatted)
429529
} else {
430530
logInfo(formatted)
431531
}
432532
}
533+
534+
/** What a span of generated Java source is, for [[forEachJavaSpan]]. */
535+
private[codegen] sealed trait JavaSpan
536+
private[codegen] object JavaSpan {
537+
/** Code: whatever is neither a literal nor a comment. */
538+
case object Code extends JavaSpan
539+
/** A string or character literal, its quotes included. */
540+
case object Literal extends JavaSpan
541+
/** A `//` comment up to its line's end, or a block comment, its delimiters included. */
542+
case object Comment extends JavaSpan
543+
}
544+
545+
/**
546+
* Calls `f` with each span of the generated Java `source` in turn: its kind, and where it
547+
* starts and ends (exclusive). The spans cover the source. A literal's backslash escapes are
548+
* honored, a `//` comment ends before its newline, and an unterminated literal or block comment
549+
* runs to the end. The one scanner of generated Java: `JdkCodeCompiler.rewriteInnerClassRefs`
550+
* rewrites the code spans, and `CodeGenerator.wholeStageSplitMarkersIn` reads the comments.
551+
*/
552+
private[codegen] def forEachJavaSpan(source: String)(f: (JavaSpan, Int, Int) => Unit): Unit = {
553+
val n = source.length
554+
var i = 0
555+
var codeStart = 0
556+
def endCode(at: Int): Unit = if (at > codeStart) f(JavaSpan.Code, codeStart, at)
557+
while (i < n) {
558+
val c = source.charAt(i)
559+
if (c == '"' || c == '\'') {
560+
endCode(i)
561+
val start = i
562+
i += 1
563+
var closed = false
564+
while (i < n && !closed) {
565+
val ch = source.charAt(i)
566+
if (ch == '\\' && i + 1 < n) {
567+
i += 2
568+
} else {
569+
i += 1
570+
closed = ch == c
571+
}
572+
}
573+
f(JavaSpan.Literal, start, i)
574+
codeStart = i
575+
} else if (c == '/' && i + 1 < n && source.charAt(i + 1) == '/') {
576+
endCode(i)
577+
val start = i
578+
i = source.indexOf('\n', i) match { case -1 => n; case end => end }
579+
f(JavaSpan.Comment, start, i)
580+
codeStart = i
581+
} else if (c == '/' && i + 1 < n && source.charAt(i + 1) == '*') {
582+
endCode(i)
583+
val start = i
584+
i = source.indexOf("*/", i + 2) match { case -1 => n; case end => end + 2 }
585+
f(JavaSpan.Comment, start, i)
586+
codeStart = i
587+
} else {
588+
i += 1
589+
}
590+
}
591+
endCode(n)
592+
}
433593
}
434594

435595
/**
@@ -472,11 +632,11 @@ object JaninoCodeCompiler extends CodeCompiler with Logging {
472632
CodeCompiler.computeByteCodeStats(evaluator.getBytecodes.asScala)
473633
} catch {
474634
case e: InternalCompilerException =>
475-
logError("Failed to compile the generated Java code.", e)
635+
CodeCompiler.logCompileFailure(e)
476636
CodeCompiler.logGeneratedCodeOnFailure(code, SQLConf.get.loggingMaxLinesForCodegen)
477637
throw QueryExecutionErrors.internalCompilerError(e)
478638
case e: CompileException =>
479-
logError("Failed to compile the generated Java code.", e)
639+
CodeCompiler.logCompileFailure(e)
480640
CodeCompiler.logGeneratedCodeOnFailure(code, SQLConf.get.loggingMaxLinesForCodegen)
481641
throw QueryExecutionErrors.compilerError(e)
482642
}
@@ -697,68 +857,19 @@ object JdkCodeCompiler extends CodeCompiler with Logging {
697857
* The rewrite is applied only to actual code spans: string literals, char
698858
* literals, and `//` / block comments are copied verbatim so that a `$Upper`
699859
* sequence inside generated string data (e.g. a column name or error message)
700-
* is never corrupted.
860+
* is never corrupted. The spans come from `CodeCompiler.forEachJavaSpan`.
701861
*/
702862
private[codegen] def rewriteInnerClassRefs(body: String, classLoader: ClassLoader): String = {
703863
val out = new java.lang.StringBuilder(body.length + 16)
704-
val code = new java.lang.StringBuilder()
705-
val n = body.length
706864
// A token's rewritten form is stable for a given classloader; memoize within
707865
// this call so repeated type references resolve at most once.
708866
val memo = mutable.HashMap.empty[String, String]
709-
710-
def flushCode(): Unit = {
711-
if (code.length > 0) {
712-
out.append(rewriteCodeSpan(code.toString, classLoader, memo))
713-
code.setLength(0)
714-
}
715-
}
716-
717-
// Copy a quoted literal (string or char) verbatim, honoring backslash escapes.
718-
def copyQuoted(start: Int, quote: Char): Int = {
719-
out.append(quote)
720-
var j = start + 1
721-
var closed = false
722-
while (j < n && !closed) {
723-
val ch = body.charAt(j)
724-
if (ch == '\\' && j + 1 < n) {
725-
out.append(ch).append(body.charAt(j + 1))
726-
j += 2
727-
} else {
728-
out.append(ch)
729-
j += 1
730-
if (ch == quote) closed = true
731-
}
732-
}
733-
j
734-
}
735-
736-
var i = 0
737-
while (i < n) {
738-
val c = body.charAt(i)
739-
if (c == '"' || c == '\'') {
740-
flushCode()
741-
i = copyQuoted(i, c)
742-
} else if (c == '/' && i + 1 < n && body.charAt(i + 1) == '/') {
743-
flushCode()
744-
while (i < n && body.charAt(i) != '\n') { out.append(body.charAt(i)); i += 1 }
745-
} else if (c == '/' && i + 1 < n && body.charAt(i + 1) == '*') {
746-
flushCode()
747-
out.append("/*")
748-
i += 2
749-
while (i < n && !(body.charAt(i) == '*' && i + 1 < n && body.charAt(i + 1) == '/')) {
750-
out.append(body.charAt(i)); i += 1
751-
}
752-
// The scan exits either at the `*/` terminator (then i + 1 < n holds by the
753-
// loop condition) or at end-of-body for an unterminated comment, whose
754-
// characters the loop already copied verbatim.
755-
if (i + 1 < n) { out.append("*/"); i += 2 }
756-
} else {
757-
code.append(c)
758-
i += 1
759-
}
867+
CodeCompiler.forEachJavaSpan(body) {
868+
case (CodeCompiler.JavaSpan.Code, start, end) =>
869+
out.append(rewriteCodeSpan(body.substring(start, end), classLoader, memo))
870+
case (_, start, end) =>
871+
out.append(body, start, end)
760872
}
761-
flushCode()
762873
out.toString
763874
}
764875

@@ -1316,9 +1427,12 @@ object JdkCodeCompiler extends CodeCompiler with Logging {
13161427
// not carry a cause, or downstream resolution can fail in non-local deployments.
13171428
val parentLoader = new ParentClassLoader(resolveLoader)
13181429

1430+
// The compile reports from the worker, so a trial compile's marker goes there with it.
1431+
val trial = CodeCompiler.activeTrial
13191432
val future = compileExecutor.submit(new Callable[(GeneratedClass, ByteCodeStats)] {
1320-
override def call(): (GeneratedClass, ByteCodeStats) =
1433+
override def call(): (GeneratedClass, ByteCodeStats) = CodeCompiler.withTrial(trial) {
13211434
doCompile(code, source, resolveLoader, parentLoader, failureLogMaxLines)
1435+
}
13221436
})
13231437
try {
13241438
// Await uninterruptibly: the result is cached and the worker must finish its
@@ -1370,7 +1484,7 @@ object JdkCodeCompiler extends CodeCompiler with Logging {
13701484
task.call().booleanValue()
13711485
} catch {
13721486
case NonFatal(e) =>
1373-
logError("Failed to compile the generated Java code.", e)
1487+
CodeCompiler.logCompileFailure(e)
13741488
logSourceOnFailure()
13751489
throw QueryExecutionErrors.internalCompilerError(
13761490
new InternalCompilerException(e.getMessage, e))
@@ -1389,7 +1503,7 @@ object JdkCodeCompiler extends CodeCompiler with Logging {
13891503
case parts => parts.mkString("\n")
13901504
}
13911505
val ex = new CompileException(message, null)
1392-
logError("Failed to compile the generated Java code.", ex)
1506+
CodeCompiler.logCompileFailure(ex)
13931507
logSourceOnFailure()
13941508
throw QueryExecutionErrors.compilerError(ex)
13951509
}

0 commit comments

Comments
 (0)