dongjoon-hyun commented on code in PR #59225:
URL: https://github.com/apache/spark/pull/59225#discussion_r4189436271
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/WholeStageCodegenExec.scala:
##########
@@ -610,6 +616,90 @@ object WholeStageCodegenExec {
// Reset generation time of Java source code.
// Visible for testing
def resetCodeGenTime(): Unit = _codeGenTime.set(0L)
+
+ /**
+ * A trial compile's outcome for `WholeStageCodegenExec.chooseCode`: the
bytecode size of each
+ * method, None where they are unknown, or the failure; and the reports and
metric updates the
+ * compile held back for the code kept.
+ */
+ private[execution] case class TrialResult(
+ methodSizes: Try[Option[Map[String, Int]]],
+ report: () => Unit)
+
+ /**
+ * Compiles `code` as a trial (`CodeCompiler.trial`). A hit in the compile
cache produces no
+ * method sizes, so the class is then dropped and compiled again. A hit
happens where the class
+ * is cached while the decision for its code is not: the decision was
evicted, or the same body
+ * was compiled earlier without being decided - with the split off, whose
code is the recorded
+ * code without its marks, or under another limit. Should another thread
load the class again in
+ * between, the second compile is a hit as well, and the sizes are unknown.
+ */
+ private[execution] def trialCompile(
+ code: CodeAndComment,
+ failureExpected: Boolean): TrialResult = {
+ def once(): (Try[Map[String, Int]], CodeCompiler.TrialCompile) = {
+ val (result, trial) =
CodeCompiler.trial(failureExpected)(CodeGenerator.compile(code))
+ (result.map(_ => trial.methodSizes), trial)
+ }
+ val (sizes, trial) = once() match {
+ case (Success(m), _) if m.isEmpty =>
+ CodeGenerator.invalidateCompiled(code)
+ once()
+ case result => result
+ }
+ TrialResult(sizes.map(m => Option.when(m.nonEmpty)(m)), () =>
trial.report())
+ }
+
+ /** The methods a class runs once, whose being interpreted costs nothing per
row. */
+ private val RunOnce = Set("<init>", "<clinit>", "init")
+
+ /** The bytecode of the methods run per row that are past `limit`. */
+ private def bytesOver(sizes: Map[String, Int], limit: Int): Int =
+ sizes.iterator.collect { case (m, size) if size > limit && !RunOnce(m) =>
size }.sum
+
+ /** What `chooseCode` decided for a stage's code; see [[splitDecisions]]. */
+ private sealed trait SplitDecision
+ /** The code in one piece runs: the JIT compiles it whole, or no split
compiles either. */
+ private case object KeepWhole extends SplitDecision
+ private case object SplitFailed extends SplitDecision
+ /** The expressions split, by their index in
`CodegenContext.wholeStageSplitsRecorded`. */
+ private case class SplitThese(indices: Option[Set[Int]]) extends
SplitDecision
+ /** Nothing decided, the method sizes being unknown: the code in one piece
runs this time. */
+ private case object Undecided extends SplitDecision
+
+ /**
+ * What `chooseCode` decided for each stage's code, by [[decisionKey]], so
that a stage generated
+ * again skips the trial compiles: the code kept in one piece, a split that
failed to compile and
+ * runs in one piece, or the expressions split. A stage met by several
threads at once is decided
+ * once, the others waiting for the decision. Sized as the compile cache is,
0 being unbounded.
+ */
+ private val splitDecisions: NonFateSharingCache[AnyRef, SplitDecision] = {
+ val builder = CacheBuilder.newBuilder()
+ val maxEntries = SQLConf.get.codegenCacheMaxEntries
+ NonFateSharingCache((if (maxEntries > 0) builder.maximumSize(maxEntries)
else builder)
Review Comment:
**[High] Passing a Guava `Cache` built in sql/core to
`NonFateSharingCache.apply` fails sql/core's Maven tests with
`NoSuchMethodError`.**
This comes from the fix for item 3 of my previous review.
`NonFateSharingCache.apply(cache: Cache[K, V])` takes a Guava type, and
spark-core's Maven jar relocates Guava (`com.google.common` to
`org.sparkproject.guava`), so the installed core jar only has
`apply(Lorg/sparkproject/guava/cache/Cache;)`; the released
spark-core_2.13-4.2.0.jar shows the same. This is the cross-module call that
SPARK-44064 added the Guava-free `apply` for (`NonFateSharingCache.scala`
L53-61), and that SPARK-50443 had to fix again for `RocksDBStateStoreProvider`.
The scheduled Maven jobs (`maven_test.yml`) run `clean install -DskipTests`
and then `-pl "$TEST_MODULES" test` without `-am`, so sql/core's unrelocated
classes run against the shaded core jar. `splitDecisions` is a strict `val` of
`object WholeStageCodegenExec` (L676-681), so the first query planned with
whole-stage codegen throws `NoSuchMethodError` from
`WholeStageCodegenExec$.<clinit>`, every later one throws
`NoClassDefFoundError`, and ScalaTest aborts the sql/core run on the
`LinkageError`. The PR's CI can't show it: `build_and_test.yml` runs the tests
with sbt, and its Maven job only installs with `-DskipTests`. Production and
the downstream modules are fine, since sql/core's own jar relocates its
references as well.
Switching only to the Guava-free `NonFateSharingCache(maxEntries, 0,
TimeUnit.SECONDS)` wouldn't be enough either: in the same runs, `decided`
(L696-702) catches the unrelocated `UncheckedExecutionException` and
`ExecutionError`, while the shaded `get` throws the relocated ones.
Could this keep a plain Guava cache in sql/core, as the previous head did,
and serialize the deciders with core's Guava-free `KeyLock`? For example,
`keyLock.withLock(key) { Option(cache.getIfPresent(key)).getOrElse { val d =
decide(); if (d != Undecided) cache.put(key, d); d } }`, which also removes the
need for `decided`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/WholeStageCodegenExec.scala:
##########
@@ -668,11 +758,181 @@ case class WholeStageCodegenExec(child: SparkPlan)(val
codegenStageId: Int)
/**
* Generates code for this subtree.
*
+ * Under `spark.sql.codegen.wholeStage.splitExpressions` the stage's code is
first generated in
+ * one piece, as with the conf off, with the code of each expression a split
could take marked,
+ * and only where there is one is that code compiled to choose which
expressions to split
+ * ([[chooseCode]]). A stage the JIT compiles whole keeps the code it had
without the split,
+ * which a split's calls would only slow down; the trial compile of the code
kept is the one
+ * [[doExecute]] finds in the compile cache.
+ *
* @return the tuple of the codegen context and the actual generated source.
*/
def doCodeGen(): (CodegenContext, CodeAndComment) = {
+ if (!conf.wholeStageSplitExpressions) {
+ generate(WholeStageSplit.Off)
+ } else if (conf.wholeStageSplitExpressionsMethodLimit == 0) {
+ generate(WholeStageSplit.Split(None))
+ } else {
+ chooseCode(generate(WholeStageSplit.Record), only =>
generate(WholeStageSplit.Split(only)),
+ WholeStageCodegenExec.trialCompile, CodeGenerator.invalidateCompiled,
+ throwSplitFailure = Utils.isTesting)
+ }
+ }
+
+ /**
+ * Which code [[doCodeGen]] returns for the stage: `recorded`, its code in
one piece with the
+ * expressions a split could take marked (`WholeStageSplit.Record`), or the
code `split` generates
+ * with some of them split. The code in one piece is compiled by `compile`
as a trial
+ * ([[WholeStageCodegenExec.trialCompile]]), and the reports of a method
past the JIT limit, with
+ * the codegen metrics, are made for the code kept.
+ *
+ * - Where it fails to compile, the split of every expression is the code
that can run, and is
+ * compiled as a trial too; where that fails as well, the code in one
piece is left to
+ * [[doExecute]], which reports its failure, falls back or throws it, as
without the split.
+ * - Otherwise the methods run for each row past the limit - the lower of
+ * `spark.sql.codegen.wholeStage.splitExpressions.methodLimit` (HotSpot's
JIT limit by default)
+ * and `spark.sql.codegen.hugeMethodLimit` - are the ones a split could
help, and the
+ * expressions marked in them are split: an expression in a method the
JIT compiles whole
+ * would only be slowed down by a split's calls. Where none is marked in
them, or the split
+ * comes out as the code in one piece, the code in one piece is kept.
+ * - The split is kept if it lowers the bytes of the methods past the
limit, and compiles;
+ * otherwise the code in one piece runs, as it did before the split
existed. A split that
+ * fails to compile is thrown where `throwSplitFailure`, which is under
testing.
+ *
+ * The decision is remembered by [[WholeStageCodegenExec.decisionKey]], so
that a stage generated
+ * again, as every execution of a query generates it, skips the trial
compiles, and threads that
+ * meet the stage at once wait for one decision; where a trial's sizes are
unknown, nothing is
+ * remembered. The variant not kept leaves the compile cache (`invalidate`),
so that it holds only
+ * code that runs and a later compile of that variant reports what it finds.
+ */
+ private[execution] def chooseCode(
+ recorded: (CodegenContext, CodeAndComment),
+ split: Option[Set[WholeStageSplitKey]] => (CodegenContext,
CodeAndComment),
+ compile: (CodeAndComment, Boolean) => WholeStageCodegenExec.TrialResult,
+ invalidate: CodeAndComment => Unit,
+ throwSplitFailure: Boolean): (CodegenContext, CodeAndComment) = {
+ import WholeStageCodegenExec._
+ val (ctx, whole) = recorded
+ if (ctx.wholeStageSplitsRecorded.isEmpty) {
+ return recorded
+ }
+ val limit = math.min(conf.wholeStageSplitExpressionsMethodLimit,
conf.hugeMethodLimit)
+ def splitOf(indices: Option[Set[Int]]) =
+ split(indices.map(_.map(ctx.wholeStageSplitsRecorded)))
+ // The code of a decision; None for indices this recording does not have.
+ def codeOf(decision: SplitDecision) = decision match {
+ case SplitThese(indices)
+ if indices.forall(_.forall(_ < ctx.wholeStageSplitsRecorded.length))
=>
+ Some(splitOf(indices))
+ case SplitThese(_) => None
+ case _ => Some(recorded)
+ }
+ val key = decisionKey(ctx, whole, limit)
+ Option(splitDecisions.getIfPresent(key)).foreach { decision =>
+ codeOf(decision) match {
+ case Some(code) => return code
+ case None => splitDecisions.invalidate(key)
+ }
+ }
+
+ // The thread that decides keeps the code it chose; one that waited
generates the code of the
+ // decision it gets.
+ var chosen: Option[(CodegenContext, CodeAndComment)] = None
+ def decide(): SplitDecision = {
+ val (decision, code) = decideFromTrials(ctx, whole, recorded, limit,
splitOf, compile,
+ invalidate, throwSplitFailure)
+ chosen = Some(code)
+ decision
+ }
+ val decision = decided(key, () => decide())
+ if (decision == Undecided) {
+ splitDecisions.invalidate(key)
+ }
+ chosen.orElse(codeOf(decision)).getOrElse(recorded)
+ }
+
+ /** The decision of `chooseCode` from the trial compiles, with the code it
keeps. */
+ private def decideFromTrials(
+ ctx: CodegenContext,
+ whole: CodeAndComment,
+ recorded: (CodegenContext, CodeAndComment),
+ limit: Int,
+ splitOf: Option[Set[Int]] => (CodegenContext, CodeAndComment),
+ compile: (CodeAndComment, Boolean) => WholeStageCodegenExec.TrialResult,
+ invalidate: CodeAndComment => Unit,
+ throwSplitFailure: Boolean): (WholeStageCodegenExec.SplitDecision,
+ (CodegenContext, CodeAndComment)) = {
+ import WholeStageCodegenExec._
+ val wholeTrial = compile(whole, true)
+ wholeTrial.methodSizes match {
+ case Failure(_) =>
+ val splitAll = splitOf(None)
+ val splitTrial = compile(splitAll._2, true)
Review Comment:
**[Low] When the code in one piece fails to compile, the split's own failure
is logged only at DEBUG, and the failing code can be compiled three times.**
This is about the fix for item 8 of my previous review. Reporting the
unsplit code's error as master does when both fail makes sense, and I agree
that not throwing an unsplit error that the split avoids is intended. The gap
is the other error: the split here is trial-compiled with `failureExpected =
true`, so when it fails as well, its error goes to DEBUG without the source,
`throwSplitFailure` isn't consulted, and `(KeepWhole, recorded)` is returned.
For a stage past 64KB, the main target, a bug in the split's code (a missing
argument, a wrong type) therefore leaves nothing above DEBUG in production,
while the `SplitFailed` branch (L918-923) does warn. Under testing, a suite
sees the unsplit "grows beyond 64 KB" error, or only `splitsCaseWhen` being
false, rather than the split's compile error. The previous head compiled this
split in `doExecute`, which logged it as an ERROR and threw it under testing.
This branch also has no `splitAll._2 == whole` check (item 6), so where no
block can be split, the first execution takes 2 generations and 3 failing
compiles, against master's 1 and 1.
The `KeepWhole` scaladoc (L662) says the code in one piece runs, while here
the stage falls back; and `SplitFailed` (L664) is handled exactly like
`KeepWhole` everywhere, so one of the two could go.
Could this skip the trial when `splitAll._2 == whole`, compile the split
with `failureExpected = !throwSplitFailure` (rethrowing under testing), and
warn like the `SplitFailed` branch does?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/WholeStageCodegenExec.scala:
##########
@@ -668,11 +758,181 @@ case class WholeStageCodegenExec(child: SparkPlan)(val
codegenStageId: Int)
/**
* Generates code for this subtree.
*
+ * Under `spark.sql.codegen.wholeStage.splitExpressions` the stage's code is
first generated in
+ * one piece, as with the conf off, with the code of each expression a split
could take marked,
+ * and only where there is one is that code compiled to choose which
expressions to split
+ * ([[chooseCode]]). A stage the JIT compiles whole keeps the code it had
without the split,
+ * which a split's calls would only slow down; the trial compile of the code
kept is the one
+ * [[doExecute]] finds in the compile cache.
+ *
* @return the tuple of the codegen context and the actual generated source.
*/
def doCodeGen(): (CodegenContext, CodeAndComment) = {
+ if (!conf.wholeStageSplitExpressions) {
+ generate(WholeStageSplit.Off)
+ } else if (conf.wholeStageSplitExpressionsMethodLimit == 0) {
+ generate(WholeStageSplit.Split(None))
+ } else {
+ chooseCode(generate(WholeStageSplit.Record), only =>
generate(WholeStageSplit.Split(only)),
+ WholeStageCodegenExec.trialCompile, CodeGenerator.invalidateCompiled,
+ throwSplitFailure = Utils.isTesting)
+ }
+ }
+
+ /**
+ * Which code [[doCodeGen]] returns for the stage: `recorded`, its code in
one piece with the
+ * expressions a split could take marked (`WholeStageSplit.Record`), or the
code `split` generates
+ * with some of them split. The code in one piece is compiled by `compile`
as a trial
+ * ([[WholeStageCodegenExec.trialCompile]]), and the reports of a method
past the JIT limit, with
+ * the codegen metrics, are made for the code kept.
+ *
+ * - Where it fails to compile, the split of every expression is the code
that can run, and is
+ * compiled as a trial too; where that fails as well, the code in one
piece is left to
+ * [[doExecute]], which reports its failure, falls back or throws it, as
without the split.
+ * - Otherwise the methods run for each row past the limit - the lower of
+ * `spark.sql.codegen.wholeStage.splitExpressions.methodLimit` (HotSpot's
JIT limit by default)
+ * and `spark.sql.codegen.hugeMethodLimit` - are the ones a split could
help, and the
+ * expressions marked in them are split: an expression in a method the
JIT compiles whole
+ * would only be slowed down by a split's calls. Where none is marked in
them, or the split
+ * comes out as the code in one piece, the code in one piece is kept.
+ * - The split is kept if it lowers the bytes of the methods past the
limit, and compiles;
+ * otherwise the code in one piece runs, as it did before the split
existed. A split that
+ * fails to compile is thrown where `throwSplitFailure`, which is under
testing.
+ *
+ * The decision is remembered by [[WholeStageCodegenExec.decisionKey]], so
that a stage generated
+ * again, as every execution of a query generates it, skips the trial
compiles, and threads that
+ * meet the stage at once wait for one decision; where a trial's sizes are
unknown, nothing is
+ * remembered. The variant not kept leaves the compile cache (`invalidate`),
so that it holds only
+ * code that runs and a later compile of that variant reports what it finds.
+ */
+ private[execution] def chooseCode(
+ recorded: (CodegenContext, CodeAndComment),
+ split: Option[Set[WholeStageSplitKey]] => (CodegenContext,
CodeAndComment),
+ compile: (CodeAndComment, Boolean) => WholeStageCodegenExec.TrialResult,
+ invalidate: CodeAndComment => Unit,
+ throwSplitFailure: Boolean): (CodegenContext, CodeAndComment) = {
+ import WholeStageCodegenExec._
+ val (ctx, whole) = recorded
+ if (ctx.wholeStageSplitsRecorded.isEmpty) {
+ return recorded
+ }
+ val limit = math.min(conf.wholeStageSplitExpressionsMethodLimit,
conf.hugeMethodLimit)
+ def splitOf(indices: Option[Set[Int]]) =
+ split(indices.map(_.map(ctx.wholeStageSplitsRecorded)))
+ // The code of a decision; None for indices this recording does not have.
+ def codeOf(decision: SplitDecision) = decision match {
+ case SplitThese(indices)
+ if indices.forall(_.forall(_ < ctx.wholeStageSplitsRecorded.length))
=>
+ Some(splitOf(indices))
+ case SplitThese(_) => None
+ case _ => Some(recorded)
+ }
+ val key = decisionKey(ctx, whole, limit)
+ Option(splitDecisions.getIfPresent(key)).foreach { decision =>
+ codeOf(decision) match {
+ case Some(code) => return code
+ case None => splitDecisions.invalidate(key)
+ }
+ }
+
+ // The thread that decides keeps the code it chose; one that waited
generates the code of the
+ // decision it gets.
+ var chosen: Option[(CodegenContext, CodeAndComment)] = None
+ def decide(): SplitDecision = {
+ val (decision, code) = decideFromTrials(ctx, whole, recorded, limit,
splitOf, compile,
+ invalidate, throwSplitFailure)
+ chosen = Some(code)
+ decision
+ }
+ val decision = decided(key, () => decide())
+ if (decision == Undecided) {
+ splitDecisions.invalidate(key)
+ }
+ chosen.orElse(codeOf(decision)).getOrElse(recorded)
+ }
+
+ /** The decision of `chooseCode` from the trial compiles, with the code it
keeps. */
+ private def decideFromTrials(
+ ctx: CodegenContext,
+ whole: CodeAndComment,
+ recorded: (CodegenContext, CodeAndComment),
+ limit: Int,
+ splitOf: Option[Set[Int]] => (CodegenContext, CodeAndComment),
+ compile: (CodeAndComment, Boolean) => WholeStageCodegenExec.TrialResult,
+ invalidate: CodeAndComment => Unit,
+ throwSplitFailure: Boolean): (WholeStageCodegenExec.SplitDecision,
+ (CodegenContext, CodeAndComment)) = {
+ import WholeStageCodegenExec._
+ val wholeTrial = compile(whole, true)
+ wholeTrial.methodSizes match {
+ case Failure(_) =>
+ val splitAll = splitOf(None)
+ val splitTrial = compile(splitAll._2, true)
+ splitTrial.methodSizes match {
+ case Success(_) =>
+ splitTrial.report()
+ (SplitThese(None), splitAll)
+ case Failure(_) =>
+ // `doExecute` compiles the code in one piece, reporting its
failure as any compile
+ // does, and falls back or throws it.
+ (KeepWhole, recorded)
+ }
+ case Success(None) =>
+ wholeTrial.report()
+ (Undecided, recorded)
+ case Success(Some(sizes)) =>
+ val over = sizes.filter { case (m, size) => size > limit &&
!RunOnce(m) }
+ val indices = if (over.isEmpty) {
+ Nil
+ } else {
+ ctx.wholeStageSplitMarkersByFunction.collect {
+ case (i, functions) if functions.exists(over.contains) => i
+ }
+ }
+ if (indices.isEmpty) {
+ wholeTrial.report()
+ (KeepWhole, recorded)
+ } else {
+ val splitCode = splitOf(Some(indices.toSet))
+ if (splitCode._2 == whole) {
Review Comment:
**[Low] This check can't fire when subexpression elimination discarded a
first pass that added helper functions, so such a stage is still compiled
twice.**
This is about the fix for item 6 of my previous review. Every `Split`
generation removes the functions of the discarded first pass, whatever `only`
holds (`CodeGenerator.scala` L2177-2179), while the `Record` code keeps them,
as the new test at `WholeStageCodegenSuite.scala` L1918-1928 pins. Their names
are fresh, so the second pass adds them again under new names. When a split
turns no block into a method, its code therefore still differs from `whole` by
those dead helpers, and it is trial-compiled again.
For example, take a recorded CASE WHEN past 8000 bytes none of whose blocks
can be split, because each block reads a column read only once, which
`ProjectExec` defers, and add to the same projection common subexpressions like
the test's `concat(cast(array(v, v + i, v * 2) AS STRING), ...)`. The discarded
first pass adds `elementToString` helpers (`ToStringBase.scala` L377-380). No
CASE WHEN block is split, but this check is false, the same over-limit method
is compiled again at L903, and the result is `KeepWhole` plus
`invalidate(splitCode._2)`: two compiles on the first execution of each key,
where the comment at L898-899, the scaladoc at `CodeGenerator.scala` L1642-1643
and your reply say one. The answers are right and the decision is remembered.
In the narrower case where a dead helper is itself past the limit, the empty
split passes the check at L905 and is remembered as `SplitThese`, so every
execution then generates the stage twice.
The test of this check (`WholeStageCodegenSuite.scala` L2095-2099) fakes
`split` to return the same object, so it doesn't go through a real generation.
Could the decision come from whether the `Split` generation turned any block
into a method, e.g. a count in `splitWholeStage` that is saved and restored
with the discarded pass like `wholeStageSplitOccurrences`, rather than from
comparing the code?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -1823,6 +2444,96 @@ object CodeGenerator extends Logging {
// method which is going to be called by the outer class instead of the many
small ones
final val MERGE_SPLIT_METHODS_THRESHOLD = 3
+ //
---------------------------------------------------------------------------------------------
+ // The whole stage split of expressions
(`CodegenContext.splitExpressionsWithSources`)
+ //
---------------------------------------------------------------------------------------------
+
+ /** A slot of a compacted mutable state array, such as
`mutableStateArray_0[3]`. */
+ private[codegen] val CompactedSlot = "([A-Za-z_$][\\w$]*)\\[\\d+\\]".r
+
+ /** A name and an opening parenthesis: a call, or a keyword such as `if`. */
+ private[codegen] val CallSite = "(?<![\\w$])([A-Za-z_$][\\w$]*)\\s*\\(".r
+
+ /**
+ * The comment marking the code of the expression of index `i` in
+ * `CodegenContext.wholeStageSplitsRecorded` under `WholeStageSplit.Record`.
+ */
+ private[sql] def wholeStageSplitMarker(i: Int): String =
s"$WholeStageSplitMarkerStart$i*/"
+
+ private val WholeStageSplitMarkerStart = "/*wholeStageSplit:"
+
+ private val WholeStageSplitMarker = """/\*wholeStageSplit:(\d+)\*/""".r
+
+ /** Whether `code` may hold a `wholeStageSplitMarker`, a check before a
scan. */
+ private[codegen] def holdsWholeStageSplitMarker(code: String): Boolean =
+ code.contains(WholeStageSplitMarkerStart)
+
+ /**
+ * Each `wholeStageSplitMarker` in the generated Java `code`, as its index
and where it starts and
+ * ends: a block comment that is exactly a marker. String and character
literals and other
+ * comments are skipped, so that text a literal holds, such as a user's
regular expression, is
+ * never taken for a marker.
+ */
+ private[sql] def wholeStageSplitMarkersIn(code: String): Seq[(Int, Int,
Int)] = {
+ val found = mutable.ArrayBuffer.empty[(Int, Int, Int)]
+ var i = 0
+ while (i < code.length) {
+ code.charAt(i) match {
+ case quote @ ('"' | '\'') =>
+ i += 1
+ while (i < code.length && code.charAt(i) != quote) {
+ if (code.charAt(i) == '\\') i += 1
+ i += 1
+ }
+ i += 1
+ case '/' if code.startsWith("//", i) =>
+ i = code.indexOf('\n', i) match { case -1 => code.length; case end
=> end + 1 }
+ case '/' if code.startsWith("/*", i) =>
+ val end = code.indexOf("*/", i + 2) match { case -1 => code.length;
case e => e + 2 }
+ code.substring(i, end) match {
+ case WholeStageSplitMarker(index) => found += ((index.toInt, i,
end))
+ case _ =>
+ }
+ i = end
+ case _ =>
+ i += 1
+ }
+ }
+ found.toSeq
+ }
+
+ /** `source` without the comments `wholeStageSplitMarker` adds
(`wholeStageSplitMarkersIn`). */
+ private[sql] def stripWholeStageSplitMarkers(source: String): String = {
+ val markers = if (holdsWholeStageSplitMarker(source)) {
+ wholeStageSplitMarkersIn(source)
+ } else {
+ Nil
+ }
+ if (markers.isEmpty) {
+ source
+ } else {
+ val stripped = new java.lang.StringBuilder(source.length)
+ var from = 0
+ markers.foreach { case (_, start, end) =>
+ stripped.append(source, from, start)
+ from = end
+ }
+ stripped.append(source, from, source.length).toString
+ }
+ }
+
+ /**
+ * `e` with the value of every lambda variable in it replaced by one value:
a fresh copy of a
+ * lambda variable gets a new value, which takes part in its equality, so
that two generations'
+ * copies of an expression holding one compare equal only without the values.
+ */
+ private[codegen] def withoutLambdaValues(e: Expression): Expression =
+ e.transformWithPruning(_.containsPattern(TreePattern.LAMBDA_VARIABLE)) {
Review Comment:
**[Medium] `withoutLambdaValues` never reaches a `NamedLambdaVariable`, so a
CASE WHEN that holds a lambda is still never split, and the new test for it has
no lambda.**
This is about the lambda half of the fix for item 7 of my previous review.
`transformWithPruning(_.containsPattern(TreePattern.LAMBDA_VARIABLE))` skips
every subtree whose pattern bits lack `LAMBDA_VARIABLE`, but
`NamedLambdaVariable` (`higherOrderFunctions.scala` L78) declares no
`nodePatterns`. Only `UnresolvedNamedLambdaVariable` (L58 there) and the
objects' `LambdaVariable` (`objects.scala` L826) set it, as
`SubexpressionEliminationSuite.scala` L626-627 also notes. With the released
4.2.0 jars, `NamedLambdaVariable(...).containsPattern(LAMBDA_VARIABLE)` and
`LambdaFunction(...).containsPattern(LAMBDA_VARIABLE)` are both false, so this
returns its input unchanged, and two generations' fresh copies of a lambda
variable still compare unequal.
For example, with `a` and `w` being `array(id, id + 1)` in `l(v, a)` and
`r(w)`, a broadcast hash join on `CASE WHEN v = 0 THEN transform(a, x -> x + 0)
... END = w`, with enough branches that the method lands between 8000 bytes and
64KB, puts the CASE WHEN in the stream side's key. The `Split` generation's key
never equals the recorded one, so `only.forall(_.contains(key))` (L1577) is
false, nothing is split, and the method stays past the JIT limit; the committed
benchmark has 3992 ms against 120 ms for 96 branches. The answers stay right.
The new test "a CASE WHEN in a join key that holds lambda variables is
split" (`WholeStageCodegenSuite.scala` L2217) builds the key from `array(CAST(v
AS DOUBLE) * k)`. `NormalizeFloatingNumbers` normalizes a `CreateArray` element
by element and adds no `transform`, so that key has no lambda, and the test
passes on the previous head too. The description ("with lambda variables'
values ignored", and the testing bullet on the `array<double>` join key) and
the scaladoc above say the same.
Pruning on `LAMBDA_FUNCTION` wouldn't work either, since the condition is
checked again at each node and the lambda's body doesn't have it. Could this be
a plain `e.transform { case v: NamedLambdaVariable => v.copy(value =
NoLambdaValue) }`, with a test whose key has a real lambda and which asserts
that the unsplit method is between 8000 and 65535 bytes before asserting the
split?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/WholeStageCodegenExec.scala:
##########
@@ -668,11 +758,181 @@ case class WholeStageCodegenExec(child: SparkPlan)(val
codegenStageId: Int)
/**
* Generates code for this subtree.
*
+ * Under `spark.sql.codegen.wholeStage.splitExpressions` the stage's code is
first generated in
+ * one piece, as with the conf off, with the code of each expression a split
could take marked,
+ * and only where there is one is that code compiled to choose which
expressions to split
+ * ([[chooseCode]]). A stage the JIT compiles whole keeps the code it had
without the split,
+ * which a split's calls would only slow down; the trial compile of the code
kept is the one
+ * [[doExecute]] finds in the compile cache.
+ *
* @return the tuple of the codegen context and the actual generated source.
*/
def doCodeGen(): (CodegenContext, CodeAndComment) = {
+ if (!conf.wholeStageSplitExpressions) {
+ generate(WholeStageSplit.Off)
+ } else if (conf.wholeStageSplitExpressionsMethodLimit == 0) {
+ generate(WholeStageSplit.Split(None))
+ } else {
+ chooseCode(generate(WholeStageSplit.Record), only =>
generate(WholeStageSplit.Split(only)),
+ WholeStageCodegenExec.trialCompile, CodeGenerator.invalidateCompiled,
+ throwSplitFailure = Utils.isTesting)
+ }
+ }
+
+ /**
+ * Which code [[doCodeGen]] returns for the stage: `recorded`, its code in
one piece with the
+ * expressions a split could take marked (`WholeStageSplit.Record`), or the
code `split` generates
+ * with some of them split. The code in one piece is compiled by `compile`
as a trial
+ * ([[WholeStageCodegenExec.trialCompile]]), and the reports of a method
past the JIT limit, with
+ * the codegen metrics, are made for the code kept.
+ *
+ * - Where it fails to compile, the split of every expression is the code
that can run, and is
+ * compiled as a trial too; where that fails as well, the code in one
piece is left to
+ * [[doExecute]], which reports its failure, falls back or throws it, as
without the split.
+ * - Otherwise the methods run for each row past the limit - the lower of
+ * `spark.sql.codegen.wholeStage.splitExpressions.methodLimit` (HotSpot's
JIT limit by default)
+ * and `spark.sql.codegen.hugeMethodLimit` - are the ones a split could
help, and the
+ * expressions marked in them are split: an expression in a method the
JIT compiles whole
+ * would only be slowed down by a split's calls. Where none is marked in
them, or the split
+ * comes out as the code in one piece, the code in one piece is kept.
+ * - The split is kept if it lowers the bytes of the methods past the
limit, and compiles;
+ * otherwise the code in one piece runs, as it did before the split
existed. A split that
+ * fails to compile is thrown where `throwSplitFailure`, which is under
testing.
+ *
+ * The decision is remembered by [[WholeStageCodegenExec.decisionKey]], so
that a stage generated
+ * again, as every execution of a query generates it, skips the trial
compiles, and threads that
+ * meet the stage at once wait for one decision; where a trial's sizes are
unknown, nothing is
+ * remembered. The variant not kept leaves the compile cache (`invalidate`),
so that it holds only
+ * code that runs and a later compile of that variant reports what it finds.
+ */
+ private[execution] def chooseCode(
+ recorded: (CodegenContext, CodeAndComment),
+ split: Option[Set[WholeStageSplitKey]] => (CodegenContext,
CodeAndComment),
+ compile: (CodeAndComment, Boolean) => WholeStageCodegenExec.TrialResult,
+ invalidate: CodeAndComment => Unit,
+ throwSplitFailure: Boolean): (CodegenContext, CodeAndComment) = {
+ import WholeStageCodegenExec._
+ val (ctx, whole) = recorded
+ if (ctx.wholeStageSplitsRecorded.isEmpty) {
+ return recorded
+ }
+ val limit = math.min(conf.wholeStageSplitExpressionsMethodLimit,
conf.hugeMethodLimit)
+ def splitOf(indices: Option[Set[Int]]) =
+ split(indices.map(_.map(ctx.wholeStageSplitsRecorded)))
+ // The code of a decision; None for indices this recording does not have.
+ def codeOf(decision: SplitDecision) = decision match {
+ case SplitThese(indices)
+ if indices.forall(_.forall(_ < ctx.wholeStageSplitsRecorded.length))
=>
+ Some(splitOf(indices))
+ case SplitThese(_) => None
+ case _ => Some(recorded)
+ }
+ val key = decisionKey(ctx, whole, limit)
+ Option(splitDecisions.getIfPresent(key)).foreach { decision =>
Review Comment:
**[Cleanup] This fast path and the out-of-range miss in `codeOf` look
unreachable now.**
Since the key holds the digest of the marked source, every index of a
remembered `SplitThese` is a marker in that source, and a recording with other
expressions has another digest, hence another key; in the deciding generation,
an index out of range would already have thrown inside the loader.
`getIfPresent` only saves the key lock on a hit, which `decided` would return
without loading anyway, and the `return` inside `foreach` (L833) is a non-local
return, which throws `NonLocalReturnControl` on every memo hit (the Databricks
Scala guide asks to avoid `return` in closures). Could `codeOf`, the fast path
and the last line become `val decision = decided(key, () => decide())`, the
`Undecided` invalidation, and `chosen.getOrElse(decision match { case
SplitThese(indices) => splitOf(indices); case _ => recorded })`?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/WholeStageCodegenSuite.scala:
##########
@@ -1799,4 +1804,808 @@ class WholeStageCodegenSuite extends SharedSparkSession
checkAnswer(query, (0 until 10).map(i => Row(i + marker)))
}
}
+
+ /** `CASE WHEN c = 1 THEN c * 1 WHEN c = 2 THEN c * 2 ... ELSE 0 END` with
`n` branches. */
+ private def largeCaseWhen(c: String, n: Int): String =
+ (1 to n).map(k => s"WHEN $c = $k THEN $c * $k").mkString("CASE ", " ", "
ELSE 0 END")
+
+ /** The largest method of `df`'s stages, compiled. */
+ private def maxStageMethodSize(df: DataFrame): Int =
+ genCode(df).map(c => CodeGenerator.compile(c)._2.maxMethodCodeSize).max
+
+ /**
+ * Whether a CASE WHEN of one of `df`'s stages was split into methods, whose
names carry the
+ * prefix of the operator they were generated for, such as
`project_caseWhen_0_0`.
+ */
+ private def splitsCaseWhen(df: DataFrame): Boolean =
+ genCode(df).exists(c => "private byte
\\w*caseWhen_\\d".r.findFirstIn(c.body).nonEmpty)
+
+ /** Whether a split method of a CASE WHEN of `df`'s stages takes a row. */
+ private def splitCaseWhenTakesRow(df: DataFrame): Boolean =
genCode(df).exists(c =>
+ "private byte \\w*caseWhen_\\d\\w*\\([^)]*InternalRow
".r.findFirstIn(c.body).nonEmpty)
+
+ /** The parameters of each split method of a CASE WHEN of `df`'s stages. */
+ private def splitCaseWhenParameters(df: DataFrame): Seq[Seq[String]] =
genCode(df).flatMap { c =>
+ "private byte
\\w*caseWhen_\\d\\w*\\(([^)]*)\\)".r.findAllMatchIn(c.body).map { m =>
+ m.group(1).split(",").map(_.trim).filter(_.nonEmpty).toSeq
+ }
+ }
+
+ /**
+ * Runs `body` with every splittable expression of a stage split: the tests
of what a split
+ * method looks like use shapes small enough for the stage to keep its code
unsplit otherwise
+ * (`spark.sql.codegen.wholeStage.splitExpressions.methodLimit`).
+ */
+ private def withSplitAlways[T](body: => T): T =
+ withSQLConf(SQLConf.WHOLESTAGE_SPLIT_EXPRESSIONS_METHOD_LIMIT.key ->
"0")(body)
+
+ /** The splits recorded in generating `df`'s stages with them recorded, as
the gate first does. */
+ private def recordedSplits(df: DataFrame): Int =
df.queryExecution.executedPlan.collect {
+ case w: WholeStageCodegenExec =>
w.generate(WholeStageSplit.Record)._1.wholeStageSplitsRecorded
+ .length
+ }.sum
+
+ /** `CASE WHEN v = 0 THEN v * 0 ... ELSE end END AS r FROM t`, `n` branches.
*/
+ private def caseWhenOverT(n: Int, end: Int): String = (0 until n)
+ .map(i => s"WHEN v = $i THEN v * $i")
+ .mkString("SELECT CASE ", " ", s" ELSE $end END AS r FROM t")
+
+ /** The rows of `query` computed without whole stage codegen, which each
test checks against. */
+ private def withoutWholeStage(query: => DataFrame): Seq[Row] =
+ withSQLConf(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false") {
+ query.collect().toSeq
+ }
+
+ test("SPARK-33301: a CASE WHEN too large for one method is split under whole
stage codegen") {
+ // Inside a stage the inputs are local variables of the stage's method,
and the branches were
+ // not split into methods at all: at 3000 branches the method went past
the JVM's 64KB limit,
+ // failed to compile, and the stage fell back to running without whole
stage codegen (under
+ // testing, the compile error is thrown). Split, with the calls to the
split methods grouped
+ // in turn, every method of the stage is small enough for HotSpot to
compile.
+ val df = spark.range(10).selectExpr(s"${largeCaseWhen("id", 3000)} AS v")
+ assert(df.queryExecution.executedPlan.exists {
+ case w: WholeStageCodegenExec => w.child.isInstanceOf[ProjectExec]
+ case _ => false
+ })
+ assert(splitsCaseWhen(df))
+ assert(maxStageMethodSize(df) <=
CodeGenerator.DEFAULT_JVM_HUGE_METHOD_LIMIT)
+ checkAnswer(df, (0L until 10L).map(i => Row(i * i)))
+ }
+
+ test("SPARK-33301: the split keeps a stage's method JIT-compilable, and can
be turned off") {
+ // Below the 64KB limit the unsplit method still goes past HotSpot's 8000
bytes, and is then
+ // never JIT-compiled.
+ def query: DataFrame = spark.range(10).selectExpr(s"${largeCaseWhen("id",
300)} AS v")
+ assert(maxStageMethodSize(query) <=
CodeGenerator.DEFAULT_JVM_HUGE_METHOD_LIMIT)
+ withSQLConf(SQLConf.WHOLESTAGE_SPLIT_EXPRESSIONS.key -> "false") {
+ assert(!splitsCaseWhen(query))
+ assert(maxStageMethodSize(query) >
CodeGenerator.DEFAULT_JVM_HUGE_METHOD_LIMIT)
+ checkAnswer(query, (0L until 10L).map(i => Row(i * i)))
+ }
+ }
+
+ test("SPARK-33301: several large CASE WHENs in one projection keep its
method small") {
+ // A projection evaluates all its outputs in one method, so the calls
every CASE WHEN leaves
+ // there add up: under whole stage codegen the calls are grouped until at
most
+ // `SPLIT_CALLS_PER_GROUP` are left per CASE WHEN, which keeps that method
JIT-compilable
+ // however many there are. Fourteen of 64 branches and four of 300.
+ for ((count, branches) <- Seq((14, 64), (4, 300))) {
+ def df: DataFrame = spark.range(100)
+ .selectExpr((1 to count).map(i => s"id + $i AS c$i"): _*)
+ .selectExpr((1 to count).map(i => s"${largeCaseWhen(s"c$i", branches)}
AS v$i"): _*)
+ assert(splitsCaseWhen(df))
+ assert(maxStageMethodSize(df) <=
CodeGenerator.DEFAULT_JVM_HUGE_METHOD_LIMIT,
+ s"$count CASE WHENs of $branches branches")
+ checkAnswer(df, withoutWholeStage(df))
+ }
+ }
+
+ test("SPARK-33301: a stage the JIT compiles whole keeps the code it had
without the split") {
+ // A CASE WHEN of 64 branches leaves the stage's method under 8000 bytes,
so the stage keeps
+ // its unsplit code, the code with the conf off, and splits only when told
to always split.
+ withTempView("t") {
+ spark.range(100).selectExpr("id AS v").createOrReplaceTempView("t")
+ val query = () => sql((0 until 64).map(i => s"WHEN v = $i THEN v * $i")
+ .mkString("SELECT CASE ", " ", " ELSE -1 END AS r FROM t"))
+ val gated = genCode(query()).map(_.body)
+ val off = withSQLConf(SQLConf.WHOLESTAGE_SPLIT_EXPRESSIONS.key ->
"false") {
+ genCode(query()).map(_.body)
+ }
+ assert(gated === off)
+ assert(!splitsCaseWhen(query()))
+ assert(withSplitAlways(splitsCaseWhen(query())))
+ checkAnswer(query(), withoutWholeStage(query()))
+ // So does a stage with no CASE WHEN whose common subexpressions take
the split pass of
+ // subexpression elimination, whose discarded first pass keeps the
helpers it added, as
+ // with the conf off.
+ val common = () => sql((0 until 8).map { i =>
+ s"concat(cast(array(v, v + $i, v * 2) AS STRING), cast(array(v, v +
$i, v * 2) AS " +
+ s"STRING)) AS c$i"
+ }.mkString("SELECT ", ", ", " FROM t"))
+ assert(genCode(common()).map(_.body) ===
+ withSQLConf(SQLConf.WHOLESTAGE_SPLIT_EXPRESSIONS.key -> "false") {
+ genCode(common()).map(_.body)
+ })
+ }
+ }
+
+ test("SPARK-33301: a stage splits only the CASE WHENs in its methods past
the limit") {
+ // The filter's CASE WHEN is generated into `processNext`, which stays
under the limit; the
+ // projection's into a method of its own, past it. Only the projection's
is split: the
+ // filter's would only be slowed down by the calls, in a method the JIT
compiles whole.
+ withTempView("t") {
+ spark.range(100).selectExpr("id AS v").createOrReplaceTempView("t")
+ def caseWhen(n: Int, base: Int): String =
+ (0 until n).map(i => s"WHEN v = $i THEN v * ${base +
i}").mkString("CASE ", " ", " END")
+ val query = () => sql(s"SELECT ${caseWhen(300, 9000)} AS r FROM t " +
+ s"WHERE ${caseWhen(64, 7000)} > 5")
+ val (ctx, _) = query().queryExecution.executedPlan.collectFirst {
+ case w: WholeStageCodegenExec => w
+ }.get.doCodeGen()
+ def methodsWith(literal: String): Set[String] =
ctx.functionsHolding(literal)
+ assert(methodsWith("9001L").nonEmpty &&
methodsWith("9001L").forall(_.contains("caseWhen")),
+ "the projection's CASE WHEN is split")
+ assert(methodsWith("7001L").nonEmpty &&
!methodsWith("7001L").exists(_.contains("caseWhen")),
+ "the filter's CASE WHEN stays where it was")
+ checkAnswer(query(), withoutWholeStage(query()))
+ }
+ }
+
+ test("SPARK-33301: the group calls of a split outside a stage are not held
to its audit") {
+ // Under testing every call a whole stage split leaves in a
subexpression's method is checked
+ // against that method's parameters. A struct hash split by
`splitExpressions` passes its own
+ // locals to its group functions, which is right, and is not checked.
+ withSQLConf(SQLConf.CODEGEN_METHOD_SPLIT_THRESHOLD.key -> "1") {
+ val struct = (0 until 200).map(i => s"'f$i', id +
$i").mkString("named_struct(", ", ", ")")
+ def df: DataFrame =
+ sql(s"SELECT hash($struct) + 1 AS a, hash($struct) + 2 AS b FROM
range(10)")
+ checkAnswer(df, withoutWholeStage(df))
+ }
+ }
+
+ test("SPARK-33301: the trial compile of a stage past the limit reports
nothing") {
+ // The unsplit code of 300 branches has a method past the JIT limit; it is
compiled only to
+ // decide to split, so the once-per-JVM warning about such a method is not
spent on it. The
+ // JDK backend compiles and reports on a thread of its own, which the
trial goes to as well.
+ withTempView("t") {
+ spark.range(100).selectExpr("id AS v").createOrReplaceTempView("t")
+ val query = caseWhenOverT(300, end = -1)
+ for (backend <- Seq("janino", "jdk")) {
+ withSQLConf(SQLConf.CODEGEN_COMPILER.key -> backend) {
+ // The warning is once per JVM and INFO after it, so the test reads
every level, of the
+ // logger the compilers report under.
+ val logs = new LogAppender("the trial compile's reports")
+ withLogAppender(logs,
+ loggerNames = Seq(classOf[CodeGenerator[_, _]].getName),
+ level = Some(Level.DEBUG)) {
+ assert(splitsCaseWhen(sql(query)))
+ sql(query).collect()
+ }
+
assert(!logs.loggingEvents.exists(_.getMessage.getFormattedMessage.contains(
+ "too long to be JIT compiled")), backend)
+ }
+ }
+ }
+ }
+
+ test("SPARK-33301: the code a stage keeps unsplit reports its method past
the JIT limit") {
+ // The trial compile holds its reports back for the code kept. Here the
default limit keeps
+ // the split of 300 branches, whose code in one piece was compiled in a
trial and is not kept.
+ // That code then leaves the compile cache, so a later compile of it
reports its method past
+ // the JIT limit, as for any compiled code: with the split turned off, and
with the limit
+ // raised past HotSpot's 8000 bytes, where the stage keeps the code in one
piece although the
+ // JIT leaves its method interpreted.
+ withTempView("t") {
+ spark.range(100).selectExpr("id AS v").createOrReplaceTempView("t")
+ val query = caseWhenOverT(300, end = -7)
+ def reports(confs: (String, String)*): Boolean = {
+ val logs = new LogAppender("the kept code's reports")
+ withLogAppender(logs,
+ loggerNames = Seq(classOf[CodeGenerator[_, _]].getName),
+ level = Some(Level.DEBUG)) {
+ // Compiled here rather than by running the query, which compiles
with the session's
+ // class loader, and so with Janino whatever the backend
(`CodeCompiler.active`).
+ withSQLConf(confs:
_*)(genCode(sql(query)).foreach(CodeGenerator.compile))
+ }
+ logs.loggingEvents.exists(_.getMessage.getFormattedMessage.contains(
+ "too long to be JIT compiled"))
+ }
+ for (backend <- Seq("janino", "jdk")) {
+ withSQLConf(SQLConf.CODEGEN_COMPILER.key -> backend) {
+ assert(splitsCaseWhen(sql(query)))
+ assert(reports(SQLConf.WHOLESTAGE_SPLIT_EXPRESSIONS.key -> "false"),
backend)
+ assert(reports(SQLConf.WHOLESTAGE_SPLIT_EXPRESSIONS_METHOD_LIMIT.key
-> "65535"), backend)
+ withSQLConf(SQLConf.WHOLESTAGE_SPLIT_EXPRESSIONS_METHOD_LIMIT.key ->
"65535") {
+ assert(!splitsCaseWhen(sql(query)))
+ }
+ }
+ }
+ }
+ }
+
+ test("SPARK-33301: which code a stage keeps, from its trial compiles") {
+ // `chooseCode`'s decisions, with the compile replaced, so that a method
can be as large as a
+ // decision needs and a split can fail to compile.
+ withTempView("t") {
+ spark.range(100).selectExpr("id AS v").createOrReplaceTempView("t")
+ val stage = sql(caseWhenOverT(300, end =
-1)).queryExecution.executedPlan.collectFirst {
+ case w: WholeStageCodegenExec => w
+ }.get
+ val recorded = stage.generate(WholeStageSplit.Record)
+ val ctx = recorded._1
+ assert(ctx.wholeStageSplitsRecorded.length === 1)
+ // The method the CASE WHEN's code lands in.
+ val Seq(method) =
ctx.wholeStageSplitMarkersByFunction.get(0).toSeq.flatten
+ val failure = new IllegalStateException("does not compile")
+ // Each case its own decision, which is remembered by the digest of the
marked code.
+ var copies = 0
+ class Run(
+ whole: Try[Option[Map[String, Int]]],
+ split: Try[Option[Map[String, Int]]],
+ splitIsWhole: Boolean = false) {
+ copies += 1
+ private val digest = s"copy $copies"
+ val code = (ctx, new CodeAndComment(recorded._2.body + s"\n/* $digest
*/",
+ recorded._2.comment))
+ val splitCode = if (splitIsWhole) {
+ code
+ } else {
+ (ctx, new CodeAndComment(code._2.body + "\n/* split */",
code._2.comment))
+ }
+ var splitAsked = Seq.empty[Option[Set[WholeStageSplitKey]]]
+ var compiles = 0
+ var invalidated = Seq.empty[CodeAndComment]
+ def choose(throwSplitFailure: Boolean = false): CodeAndComment = {
+ ctx.wholeStageSplitDigest = digest
+ stage.chooseCode(code,
+ only => { splitAsked :+= only; splitCode },
+ (c, _) => {
+ compiles += 1
+ WholeStageCodegenExec.TrialResult(if (c eq code._2) whole else
split, () => ())
+ },
+ c => invalidated :+= c, throwSplitFailure)._2
+ }
+ }
+ def sizes(m: (String, Int)*): Try[Option[Map[String, Int]]] =
Success(Some(m.toMap))
+ val notCompiling = Failure(failure)
+ val unknown = Success(None)
+ val recordedOnly = Some(Set(ctx.wholeStageSplitsRecorded.head))
+ // The split is kept where it lowers the bytes past the limit, and the
code in one piece
+ // leaves the compile cache.
+ val kept = new Run(sizes(method -> 20000), sizes(method -> 3000))
+ assert(kept.choose() eq kept.splitCode._2)
+ assert(kept.splitAsked === Seq(recordedOnly) && kept.invalidated ===
Seq(kept.code._2))
+ // A larger method that holds nothing marked does not stop the split of
the one that does.
+ val beside = new Run(sizes(method -> 9000, "other" -> 12000),
+ sizes(method -> 1500, "other" -> 12000))
+ assert(beside.choose() eq beside.splitCode._2)
+ // A method past the limit that holds nothing marked, or a method run
once, is left alone,
+ // with no split generated.
+ for (whole <- Seq(sizes(method -> 1000, "other" -> 20000),
+ sizes(method -> 1000, "init" -> 20000))) {
+ val left = new Run(whole, notCompiling)
+ assert(left.choose() eq left.code._2)
+ assert(left.splitAsked.isEmpty && left.compiles === 1)
+ }
+ // A split that leaves the bytes past the limit as they were is not
kept, and leaves the
+ // compile cache itself.
+ val same = new Run(sizes(method -> 20000), sizes(method -> 20000))
+ assert(same.choose() eq same.code._2)
+ assert(same.invalidated === Seq(same.splitCode._2))
+ // A split that comes out as the code in one piece is neither compiled
nor dropped: the code
+ // in one piece is kept, and its class stays cached for `doExecute`.
+ val identical = new Run(sizes(method -> 20000), notCompiling,
splitIsWhole = true)
+ assert(identical.choose() eq identical.code._2)
+ assert(identical.compiles === 1 && identical.invalidated.isEmpty)
+ // A split that fails to compile leaves the code in one piece, which
compiled, and is not
+ // tried again; under testing the failure is thrown.
+ val fails = new Run(sizes(method -> 20000), notCompiling)
+ assert(fails.choose() eq fails.code._2)
+ assert(fails.choose() eq fails.code._2)
+ assert(fails.compiles === 2)
+ val thrown = new Run(sizes(method -> 20000), notCompiling)
+ assert(intercept[IllegalStateException](thrown.choose(throwSplitFailure
= true)) eq failure)
+ // Code in one piece that fails to compile has every expression split
where that compiles,
+ // and the decision is remembered: the next generation asks for the same
split with no trial.
+ val broken = new Run(notCompiling, sizes(method -> 3000))
+ assert(broken.choose() eq broken.splitCode._2)
+ assert(broken.choose() eq broken.splitCode._2)
+ assert(broken.splitAsked === Seq(None, None) && broken.compiles === 2)
+ // Where the split fails too, the code in one piece is left to
`doExecute`, which reports
+ // its failure as without the split, also under testing; that is
remembered as well.
+ val bothFail = new Run(notCompiling, notCompiling)
+ assert(bothFail.choose(throwSplitFailure = true) eq bothFail.code._2)
+ assert(bothFail.choose(throwSplitFailure = true) eq bothFail.code._2)
+ assert(bothFail.compiles === 2)
+ // Sizes unknown, a second compile cache hit, keep the code in one piece
and decide nothing,
+ // so the next generation compiles again.
+ val unknownWhole = new Run(unknown, notCompiling)
+ assert(unknownWhole.choose() eq unknownWhole.code._2)
+ assert(unknownWhole.choose() eq unknownWhole.code._2)
+ assert(unknownWhole.compiles === 2 && unknownWhole.splitAsked.isEmpty)
+ val unknownSplit = new Run(sizes(method -> 20000), unknown)
+ assert(unknownSplit.choose() eq unknownSplit.code._2)
+ assert(unknownSplit.choose() eq unknownSplit.code._2)
+ assert(unknownSplit.compiles === 4)
+ }
+ }
+
+ test("SPARK-33301: a decision is not replayed under another
methodSplitThreshold") {
+ // Which CASE WHENs are recorded depends on the threshold, while the code
without the marks
+ // does not, so a decision made under one threshold names other
expressions under another.
+ // Each order of the two, in one session, runs the query right.
+ withTempView("t") {
+ spark.range(10).selectExpr("id").createOrReplaceTempView("t")
+ val a = (0 until 60).map(k => s"WHEN id = $k THEN id *
$k").mkString("CASE ", " ", " END")
+ val b = (0 until 300).map(k => s"WHEN id + 1 = $k THEN id * $k")
+ .mkString("CASE ", " ", " END")
+ for ((first, second) <- Seq(("1024", "40000"), ("40000", "1024"))) {
+ val query = s"SELECT $a AS a, $b AS b, $first AS threshold FROM t"
+ for (threshold <- Seq(first, second)) {
+ withSQLConf(SQLConf.CODEGEN_METHOD_SPLIT_THRESHOLD.key -> threshold)
{
+ checkAnswer(sql(query), withoutWholeStage(sql(query)))
+ }
+ }
+ }
+ }
+ }
+
+ test("SPARK-33301: a literal holding the text of a mark is left as it is") {
+ // A foldable pattern is written into the code as a Java literal. The
stage records its CASE
+ // WHEN and keeps its code in one piece, the marks stripped, and the
literal with them.
+ withTempView("t") {
+ spark.range(100).selectExpr("id AS v", "CAST(id AS STRING) AS s")
+ .createOrReplaceTempView("t")
+ val caseWhen = (0 until 64).map(k => s"WHEN v = $k THEN v * $k")
+ .mkString("CASE ", " ", " ELSE -1 END")
+ def df: DataFrame = sql(s"SELECT $caseWhen AS r, s RLIKE
'/*wholeStageSplit:0*/' AS m " +
+ "FROM t")
+ assert(recordedSplits(df) > 0 && !splitsCaseWhen(df))
+ checkAnswer(df, withoutWholeStage(df))
+ checkAnswer(sql(s"SELECT count(*) FROM t WHERE s RLIKE
'/*wholeStageSplit:0*/' AND " +
+ s"$caseWhen > -2"), Row(0))
+ }
+ val code = "a(); /*wholeStageSplit:0*/ b(\"/*wholeStageSplit:1*/\", '\\'')
// x\n" +
+ "/* other */ c(); /*wholeStageSplit:2*/"
+ assert(CodeGenerator.stripWholeStageSplitMarkers(code) ===
+ "a(); b(\"/*wholeStageSplit:1*/\", '\\'') // x\n/* other */ c(); ")
+ }
+
+ test("SPARK-33301: a simple CASE over a collated column is split") {
+ // The collated column is in every condition, so it is a common
subexpression, whose code is
+ // the column's: the split methods take the column.
+ withTempView("t") {
+ spark.range(100).selectExpr("CAST(id % 20 AS STRING) AS
s").createOrReplaceTempView("t")
Review Comment:
**[Test] The data never matches a branch, so `checkAnswer` doesn't check any
value computed in a split method.**
`CAST(id % 20 AS STRING)` is "0" to "19", and the branches compare it with
"k0" to "k299" under `UNICODE_CI`, which never match: with the released
spark-unsafe 4.2.0 and icu4j, none of the 6000 pairs is equal, while "K1" and
"k1" are. So all 100 rows take `ELSE 0`, and a split method returning a wrong
value or a wrong `isNull` for a matching branch would still pass; only a
compile failure or an exception would show. Could the data be e.g. `concat('K',
CAST(id % 20 AS STRING)) AS s`? The pass-through of an identity `Cast` in
`generatedChild` (item 4 of my previous review) has no end-to-end test either;
a cast that only renames struct fields would cover it.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -210,9 +339,10 @@ class CodegenContext extends Logging {
/**
* The slots a `CommonExpressionRef` reads: the value and its nullness, plus
the flag saying
- * whether this row has computed them yet, and the definition to compute
them from.
+ * whether this row has computed them yet, and the definition, with its id,
to compute them from.
*/
case class CommonExprSlots(
+ id: Long,
Review Comment:
**[Nit] `id` is never read.**
`currentCommonExprs` is keyed by the id (L520-525), and nothing reads `.id`
from a `CommonExprSlots`, so the field only takes part in `equals`, `hashCode`
and `copy`. Could it go?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -1271,49 +1759,159 @@ class CodegenContext extends Logging {
DEFAULT_JVM_HUGE_METHOD_LIMIT / 2 / (MAX_BYTES_PER_SPLIT_CALL + 2 *
arguments.length))
/**
- * The calls to the functions `splitExpressions` split out, grouped into
functions of their own
- * where there are too many for the calling method to be JIT-compiled, level
by level.
+ * The calls to the functions `splitExpressions` split out, grouped by
`groupSplitCalls` where
+ * there are too many for the calling method to be JIT-compiled. Every call
takes `arguments`, so
+ * the groups are of `SPLIT_CALLS_PER_GROUP` calls, and the calls stay in
the calling method while
+ * they fit in `splitCallsPerMethod`.
+ */
+ private def groupFunctionCalls(
+ calls: Seq[String],
+ funcName: String,
+ arguments: Seq[(String, String)],
+ returnType: String,
+ makeSplitFunction: String => String,
+ foldFunctions: Seq[String] => String): Seq[String] = {
+ groupSplitCalls(calls.map(SplitCall(_, arguments)), funcName, returnType,
makeSplitFunction,
+ foldFunctions, callsLeft = splitCallsPerMethod(arguments)).map(_.call)
+ }
+
+ /** A call to a split function, and the (type, name) arguments it passes. */
+ private case class SplitCall(call: String, arguments: Seq[(String, String)])
+
+ /**
+ * Adds the function `name`, taking `arguments` and running `body`, and
returns a call to it.
+ * Where `audited`, its arguments are kept under testing for
`assertSplitCallsWithin`: those of a
+ * whole stage split, computed from what its code reads where it was
generated, which an operator
+ * moving that code into a method of its own must take. The other splits'
arguments are their
+ * callers' own, which may be locals of the method they are made in.
+ */
+ private def splitCall(
+ name: String,
+ arguments: Seq[(String, String)],
+ returnType: String,
+ body: String,
+ audited: Boolean = false): SplitCall = {
+ if (Utils.isTesting) {
+ // Passing global variables to the split method is dangerous, as any
mutating to it is
+ // ignored and may lead to unexpected behavior.
+ arguments.foreach { case (_, argument) =>
+ assert(!mutableStateNames.contains(argument),
+ s"split function argument $argument cannot be a global variable.")
+ }
+ }
+ val argDefinitionString = arguments.map { case (t, argument) => s"$t
$argument" }.mkString(", ")
+ val code =
+ s"""
+ |private $returnType $name($argDefinitionString) {
+ | $body
+ |}
+ """.stripMargin
+ val call = s"${addNewFunction(name,
code)}(${arguments.map(_._2).mkString(", ")})"
+ if (INPUT_ROW != null && arguments.exists(_._2 == INPUT_ROW)) {
+ functionsTakingRow += name
+ }
+ if (audited && Utils.isTesting) {
+ splitCallArguments(name) = arguments.map(_._2)
+ }
+ SplitCall(call, arguments)
+ }
+
+ /**
+ * The calls to split functions, grouped into functions of their own, level
by level, until at
+ * most `callsLeft` remain in the calling method.
*
* The calls are what is left in the calling method, one per block of code,
and a `CASE WHEN`
* of a thousand branches has hundreds of blocks: their calls alone take the
caller past the
* 8000 bytes HotSpot compiles
(`CodeGenerator.DEFAULT_JVM_HUGE_METHOD_LIMIT`), so the method
* every row goes through runs interpreted. A call and its fold compile to
at most about 20 bytes
* plus 2 for each argument - the invocation, a load per argument, and the
fold's store, test
- * and branch - so the calls are left as they are while they fit in half of
that limit, the rest
- * being the caller's own code, and grouped otherwise. The groups are small,
of
- * `CodeGenerator.SPLIT_CALLS_PER_GROUP` calls, so that C2 can inline the
split functions into
- * each group within its inlining budget; a group of hundreds of calls stays
compilable but has
- * most of its calls left out of inlining, and runs slower. A group is
folded with
- * `foldFunctions` and wrapped with `makeSplitFunction`, which is what
- * `generateInnerClassesFunctionCalls` does with the functions of one inner
class, so every
- * caller of `splitExpressions` supports it.
- */
- private def groupFunctionCalls(
- calls: Seq[String],
+ * and branch - so outside whole stage codegen the calls are left as they
are while they fit in
+ * half of that limit, the rest being the caller's own code
(`splitCallsPerMethod`). Under whole
+ * stage codegen the caller is the operator's method, which holds every
expression's calls, so
+ * at most `SPLIT_CALLS_PER_GROUP` are left there for each.
+ *
+ * The groups are small, of `CodeGenerator.SPLIT_CALLS_PER_GROUP` calls, so
that C2 can inline
+ * the split functions into each group within its inlining budget; a group
of hundreds of calls
+ * stays compilable but has most of its calls left out of inlining, and runs
slower. A group
+ * takes the union of what its calls take, and is closed when one more call
would take it past
+ * those calls or past the parameters a method can have; where no group can
hold two calls, the
+ * calls are left as they are. A group is folded with `foldFunctions` and
wrapped with
+ * `makeSplitFunction`, which is what `generateInnerClassesFunctionCalls`
does with the functions
+ * of one inner class, so every caller of `splitExpressions` supports it.
+ */
+ private def groupSplitCalls(
+ calls: Seq[SplitCall],
funcName: String,
- arguments: Seq[(String, String)],
returnType: String,
makeSplitFunction: String => String,
foldFunctions: Seq[String] => String,
- level: Int = 0): Seq[String] = {
- val callsPerMethod = splitCallsPerMethod(arguments)
- if (calls.length <= callsPerMethod) {
+ callsLeft: Int,
+ level: Int = 0,
+ audited: Boolean = false): Seq[SplitCall] = {
+ if (calls.length <= callsLeft) {
calls
} else {
- val argDefinitionString = arguments.map { case (t, name) => s"$t $name"
}.mkString(", ")
- val argInvocationString = arguments.map(_._2).mkString(", ")
- val groupCalls = calls.grouped(SPLIT_CALLS_PER_GROUP).zipWithIndex.map {
case (group, i) =>
- val name = s"${funcName}_group${level}_$i"
- val code =
- s"""
- |private $returnType $name($argDefinitionString) {
- | ${makeSplitFunction(foldFunctions(group))}
- |}
- """.stripMargin
- s"${addNewFunction(name, code)}($argInvocationString)"
- }.toSeq
- groupFunctionCalls(groupCalls, funcName, arguments, returnType,
makeSplitFunction,
- foldFunctions, level + 1)
+ val groups = mutable.ArrayBuffer.empty[mutable.ArrayBuffer[SplitCall]]
+ var taken = Seq.empty[(String, String)]
+ calls.foreach { c =>
+ val union = (taken ++ c.arguments).distinctBy(_._2)
+ if (groups.isEmpty || groups.last.length == SPLIT_CALLS_PER_GROUP ||
+ !isValidParamLength(paramSlots(union))) {
+ groups += mutable.ArrayBuffer(c)
+ taken = c.arguments
+ } else {
+ groups.last += c
+ taken = union
+ }
+ }
+ if (groups.length == calls.length) {
+ calls
+ } else {
+ val grouped = groups.zipWithIndex.map { case (group, i) =>
+ splitCall(s"${funcName}_group${level}_$i",
+ group.flatMap(_.arguments).distinctBy(_._2).toSeq, returnType,
+ makeSplitFunction(foldFunctions(group.map(_.call).toSeq)), audited)
+ }.toSeq
+ groupSplitCalls(grouped, funcName, returnType, makeSplitFunction,
foldFunctions,
+ callsLeft, level + 1, audited)
+ }
+ }
+ }
+
+ /** The parameter slots of a method taking `arguments`, `this` included. */
+ private def paramSlots(arguments: Seq[(String, String)]): Int =
Review Comment:
**[Cleanup] `paramSlots` re-implements `calculateParamLengthFromExprValues`,
and the `functionsTakingRow` registration is written twice.**
At L1674-1675 the `VariableValue`s are turned into (type name, name) pairs,
and this method derives the slot count back from the type name, while
`CodeGenerator.calculateParamLengthFromExprValues` (L3108) counts the values
themselves; only the merged argument lists at L1859 need the string form.
Separately, the rule `if (INPUT_ROW != null && <arguments>.exists(<name> ==
INPUT_ROW)) functionsTakingRow += <function>` is copied in
`CommonExprSlots.build` (L429-431) and `splitCall` (L1810-1812); one helper
that adds the function and registers it would keep a third method builder from
missing it.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -1148,6 +1291,145 @@ class CodegenContext extends Logging {
}
}
+ /**
+ * What the code generated for `roots` reads from the scope it was generated
in, for moving that
+ * code into a method: the local variables to pass and the inputs to
evaluate first, or None
+ * where `policy` refuses the method. The code reads its inputs from
`currentVars` and, where
+ * the operator sets one, from `INPUT_ROW`; from the states in `subExprs`;
and, inside a `With`,
+ * from the definitions its references compute. A field or a literal is read
as it stands.
+ *
+ * Every place that moves generated code into a method walks the code's
expressions with this,
+ * so that a node binding a variable needs one change here; where those
places need different
+ * answers, `policy` says which (see [[InputPolicy]]).
+ */
+ private[sql] def collectInputs(
+ roots: Seq[Expression],
+ policy: InputPolicy,
+ subExprs: Map[ExpressionEquals, SubExprEliminationState]):
Option[CollectedInputs] = {
+ val args = mutable.LinkedHashMap.empty[String, VariableValue]
+ val inputsToEvaluate = mutable.ArrayBuffer.empty[ExprCode]
+ // Whether the code can read `v` from the method: false for a value no
parameter can carry
+ // where the policy asks for names. `ExpandExec` hands out a
`VariableValue` naming a slot of a
+ // compacted mutable state array, and a `SimpleExprValue` is an expression
rather than a name
+ // -- `posexplode_outer` gives its position the nullness `index == -1`,
and a `Byte` or `Short`
+ // literal's value is `(byte)1`. A field or a literal needs no parameter
and is read as it
+ // stands; so is a field handed out as a `VariableValue` where the policy
says so, passing one
+ // being refused by the split methods' own check, a slot of a compacted
array included.
+ def take(v: ExprValue): Boolean = v match {
+ case local: VariableValue =>
+ val name = local.variableName
+ val isName = name.nonEmpty &&
Character.isJavaIdentifierStart(name.head) &&
+ name.forall(Character.isJavaIdentifierPart)
+ // A slot of a compacted array is a field, since no parameter can be
named `array[i]`,
+ // where the whole stage split is on; with it off, every policy
answers as it did before
+ // that split existed. Another field is passed where the policy says
so.
+ val isField = (isCompactedSlot(name) &&
Review Comment:
**[Low] With the conf on, which is the default, this rule applies to every
stage, so "keeps exactly the code it had before this PR" holds only with the
conf off; and the UNPIVOT fix it brings isn't mentioned or tested.**
This is about the fix for item 9 of my previous review, which made the
conf-off case exact. With the conf on, `doCodeGen` generates every stage in
`Record` mode (`WholeStageCodegenExec.scala` L771-776), so `wholeStageSplit !=
WholeStageSplit.Off` holds in every stage, with or without a CASE WHEN, and
`collectInputs` takes a compacted slot as a field under every policy wherever
an input is one (e.g. `ExpandExec`'s non-primitive outputs).
That is a visible change, and a good one. On the released 4.3.0, i.e.
master's behavior, `SELECT c, f AS x, concat(f, 'z') AS y FROM t UNPIVOT (v FOR
c IN (a, b))`, with `f` = `concat(upper(v), lower(v), reverse(v), trim(v),
ltrim(v), rtrim(v), initcap(v), repeat(v, 2), lpad(v, 10, 'x'), rpad(v, 10,
'y'))`, generates `private void
project_subExpr_9(org.apache.spark.unsafe.types.UTF8String
expand_mutableStateArray_0[1], boolean expand_resultIsNull_1)`, which fails
with "')' expected instead of '['" and falls back. With this PR and the conf
on, the slot is read as a field and the stage compiles (from reading the code);
with the conf off it fails again, with no CASE WHEN anywhere. Could the
description mention it, and a test pin it?
A few related docs:
- The `doCodeGen` scaladoc (`WholeStageCodegenExec.scala` L761-765) says the
code is "first generated in one piece, as with the conf off" and that a stage
the JIT compiles whole "keeps the code it had without the split".
- The `splitExpressions` doc (`SQLConf.scala` L3078-3079) and the comment at
L2136-2140 say that the functions of subexpression elimination's discarded
first pass are removed with the split on, but only a `Split` generation removes
them (L2177-2179), so every stage kept in one piece keeps them, as the test at
`WholeStageCodegenSuite.scala` L1918-1928 pins. The description's "The
discarded pass" paragraph and testing bullet say "removed", while its gate
paragraph says a stage kept whole keeps "the helper functions of a discarded
subexpression pass".
- `policy.fieldsNeedNoArgument ||` above never decides: only the whole-stage
policy sets it (L183), and that policy only runs with the split on, so the
condition can be `isCompactedSlot(name) && wholeStageSplit !=
WholeStageSplit.Off`, as the comment says.
- "The fork's full CI passed on `484b4146d1b`" in the description is stale:
that is the 7th of the 13 commits, and there's no green run on b4ddb07889c yet.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/WholeStageCodegenSuite.scala:
##########
@@ -1799,4 +1804,808 @@ class WholeStageCodegenSuite extends SharedSparkSession
checkAnswer(query, (0 until 10).map(i => Row(i + marker)))
}
}
+
+ /** `CASE WHEN c = 1 THEN c * 1 WHEN c = 2 THEN c * 2 ... ELSE 0 END` with
`n` branches. */
+ private def largeCaseWhen(c: String, n: Int): String =
+ (1 to n).map(k => s"WHEN $c = $k THEN $c * $k").mkString("CASE ", " ", "
ELSE 0 END")
+
+ /** The largest method of `df`'s stages, compiled. */
+ private def maxStageMethodSize(df: DataFrame): Int =
+ genCode(df).map(c => CodeGenerator.compile(c)._2.maxMethodCodeSize).max
+
+ /**
+ * Whether a CASE WHEN of one of `df`'s stages was split into methods, whose
names carry the
+ * prefix of the operator they were generated for, such as
`project_caseWhen_0_0`.
+ */
+ private def splitsCaseWhen(df: DataFrame): Boolean =
+ genCode(df).exists(c => "private byte
\\w*caseWhen_\\d".r.findFirstIn(c.body).nonEmpty)
+
+ /** Whether a split method of a CASE WHEN of `df`'s stages takes a row. */
+ private def splitCaseWhenTakesRow(df: DataFrame): Boolean =
genCode(df).exists(c =>
+ "private byte \\w*caseWhen_\\d\\w*\\([^)]*InternalRow
".r.findFirstIn(c.body).nonEmpty)
+
+ /** The parameters of each split method of a CASE WHEN of `df`'s stages. */
+ private def splitCaseWhenParameters(df: DataFrame): Seq[Seq[String]] =
genCode(df).flatMap { c =>
+ "private byte
\\w*caseWhen_\\d\\w*\\(([^)]*)\\)".r.findAllMatchIn(c.body).map { m =>
+ m.group(1).split(",").map(_.trim).filter(_.nonEmpty).toSeq
+ }
+ }
+
+ /**
+ * Runs `body` with every splittable expression of a stage split: the tests
of what a split
+ * method looks like use shapes small enough for the stage to keep its code
unsplit otherwise
+ * (`spark.sql.codegen.wholeStage.splitExpressions.methodLimit`).
+ */
+ private def withSplitAlways[T](body: => T): T =
+ withSQLConf(SQLConf.WHOLESTAGE_SPLIT_EXPRESSIONS_METHOD_LIMIT.key ->
"0")(body)
+
+ /** The splits recorded in generating `df`'s stages with them recorded, as
the gate first does. */
+ private def recordedSplits(df: DataFrame): Int =
df.queryExecution.executedPlan.collect {
+ case w: WholeStageCodegenExec =>
w.generate(WholeStageSplit.Record)._1.wholeStageSplitsRecorded
+ .length
+ }.sum
+
+ /** `CASE WHEN v = 0 THEN v * 0 ... ELSE end END AS r FROM t`, `n` branches.
*/
+ private def caseWhenOverT(n: Int, end: Int): String = (0 until n)
+ .map(i => s"WHEN v = $i THEN v * $i")
+ .mkString("SELECT CASE ", " ", s" ELSE $end END AS r FROM t")
+
+ /** The rows of `query` computed without whole stage codegen, which each
test checks against. */
+ private def withoutWholeStage(query: => DataFrame): Seq[Row] =
+ withSQLConf(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false") {
+ query.collect().toSeq
+ }
+
+ test("SPARK-33301: a CASE WHEN too large for one method is split under whole
stage codegen") {
+ // Inside a stage the inputs are local variables of the stage's method,
and the branches were
+ // not split into methods at all: at 3000 branches the method went past
the JVM's 64KB limit,
+ // failed to compile, and the stage fell back to running without whole
stage codegen (under
+ // testing, the compile error is thrown). Split, with the calls to the
split methods grouped
+ // in turn, every method of the stage is small enough for HotSpot to
compile.
+ val df = spark.range(10).selectExpr(s"${largeCaseWhen("id", 3000)} AS v")
+ assert(df.queryExecution.executedPlan.exists {
+ case w: WholeStageCodegenExec => w.child.isInstanceOf[ProjectExec]
+ case _ => false
+ })
+ assert(splitsCaseWhen(df))
+ assert(maxStageMethodSize(df) <=
CodeGenerator.DEFAULT_JVM_HUGE_METHOD_LIMIT)
+ checkAnswer(df, (0L until 10L).map(i => Row(i * i)))
+ }
+
+ test("SPARK-33301: the split keeps a stage's method JIT-compilable, and can
be turned off") {
+ // Below the 64KB limit the unsplit method still goes past HotSpot's 8000
bytes, and is then
+ // never JIT-compiled.
+ def query: DataFrame = spark.range(10).selectExpr(s"${largeCaseWhen("id",
300)} AS v")
+ assert(maxStageMethodSize(query) <=
CodeGenerator.DEFAULT_JVM_HUGE_METHOD_LIMIT)
+ withSQLConf(SQLConf.WHOLESTAGE_SPLIT_EXPRESSIONS.key -> "false") {
+ assert(!splitsCaseWhen(query))
+ assert(maxStageMethodSize(query) >
CodeGenerator.DEFAULT_JVM_HUGE_METHOD_LIMIT)
+ checkAnswer(query, (0L until 10L).map(i => Row(i * i)))
+ }
+ }
+
+ test("SPARK-33301: several large CASE WHENs in one projection keep its
method small") {
+ // A projection evaluates all its outputs in one method, so the calls
every CASE WHEN leaves
+ // there add up: under whole stage codegen the calls are grouped until at
most
+ // `SPLIT_CALLS_PER_GROUP` are left per CASE WHEN, which keeps that method
JIT-compilable
+ // however many there are. Fourteen of 64 branches and four of 300.
+ for ((count, branches) <- Seq((14, 64), (4, 300))) {
+ def df: DataFrame = spark.range(100)
+ .selectExpr((1 to count).map(i => s"id + $i AS c$i"): _*)
+ .selectExpr((1 to count).map(i => s"${largeCaseWhen(s"c$i", branches)}
AS v$i"): _*)
+ assert(splitsCaseWhen(df))
+ assert(maxStageMethodSize(df) <=
CodeGenerator.DEFAULT_JVM_HUGE_METHOD_LIMIT,
+ s"$count CASE WHENs of $branches branches")
+ checkAnswer(df, withoutWholeStage(df))
+ }
+ }
+
+ test("SPARK-33301: a stage the JIT compiles whole keeps the code it had
without the split") {
+ // A CASE WHEN of 64 branches leaves the stage's method under 8000 bytes,
so the stage keeps
+ // its unsplit code, the code with the conf off, and splits only when told
to always split.
+ withTempView("t") {
+ spark.range(100).selectExpr("id AS v").createOrReplaceTempView("t")
+ val query = () => sql((0 until 64).map(i => s"WHEN v = $i THEN v * $i")
+ .mkString("SELECT CASE ", " ", " ELSE -1 END AS r FROM t"))
+ val gated = genCode(query()).map(_.body)
+ val off = withSQLConf(SQLConf.WHOLESTAGE_SPLIT_EXPRESSIONS.key ->
"false") {
+ genCode(query()).map(_.body)
+ }
+ assert(gated === off)
+ assert(!splitsCaseWhen(query()))
+ assert(withSplitAlways(splitsCaseWhen(query())))
+ checkAnswer(query(), withoutWholeStage(query()))
+ // So does a stage with no CASE WHEN whose common subexpressions take
the split pass of
+ // subexpression elimination, whose discarded first pass keeps the
helpers it added, as
+ // with the conf off.
+ val common = () => sql((0 until 8).map { i =>
+ s"concat(cast(array(v, v + $i, v * 2) AS STRING), cast(array(v, v +
$i, v * 2) AS " +
+ s"STRING)) AS c$i"
+ }.mkString("SELECT ", ", ", " FROM t"))
+ assert(genCode(common()).map(_.body) ===
+ withSQLConf(SQLConf.WHOLESTAGE_SPLIT_EXPRESSIONS.key -> "false") {
+ genCode(common()).map(_.body)
+ })
+ }
+ }
+
+ test("SPARK-33301: a stage splits only the CASE WHENs in its methods past
the limit") {
+ // The filter's CASE WHEN is generated into `processNext`, which stays
under the limit; the
+ // projection's into a method of its own, past it. Only the projection's
is split: the
+ // filter's would only be slowed down by the calls, in a method the JIT
compiles whole.
+ withTempView("t") {
+ spark.range(100).selectExpr("id AS v").createOrReplaceTempView("t")
+ def caseWhen(n: Int, base: Int): String =
+ (0 until n).map(i => s"WHEN v = $i THEN v * ${base +
i}").mkString("CASE ", " ", " END")
+ val query = () => sql(s"SELECT ${caseWhen(300, 9000)} AS r FROM t " +
+ s"WHERE ${caseWhen(64, 7000)} > 5")
+ val (ctx, _) = query().queryExecution.executedPlan.collectFirst {
+ case w: WholeStageCodegenExec => w
+ }.get.doCodeGen()
+ def methodsWith(literal: String): Set[String] =
ctx.functionsHolding(literal)
+ assert(methodsWith("9001L").nonEmpty &&
methodsWith("9001L").forall(_.contains("caseWhen")),
+ "the projection's CASE WHEN is split")
+ assert(methodsWith("7001L").nonEmpty &&
!methodsWith("7001L").exists(_.contains("caseWhen")),
+ "the filter's CASE WHEN stays where it was")
+ checkAnswer(query(), withoutWholeStage(query()))
+ }
+ }
+
+ test("SPARK-33301: the group calls of a split outside a stage are not held
to its audit") {
+ // Under testing every call a whole stage split leaves in a
subexpression's method is checked
+ // against that method's parameters. A struct hash split by
`splitExpressions` passes its own
+ // locals to its group functions, which is right, and is not checked.
+ withSQLConf(SQLConf.CODEGEN_METHOD_SPLIT_THRESHOLD.key -> "1") {
+ val struct = (0 until 200).map(i => s"'f$i', id +
$i").mkString("named_struct(", ", ", ")")
+ def df: DataFrame =
+ sql(s"SELECT hash($struct) + 1 AS a, hash($struct) + 2 AS b FROM
range(10)")
+ checkAnswer(df, withoutWholeStage(df))
+ }
+ }
+
+ test("SPARK-33301: the trial compile of a stage past the limit reports
nothing") {
+ // The unsplit code of 300 branches has a method past the JIT limit; it is
compiled only to
+ // decide to split, so the once-per-JVM warning about such a method is not
spent on it. The
+ // JDK backend compiles and reports on a thread of its own, which the
trial goes to as well.
+ withTempView("t") {
+ spark.range(100).selectExpr("id AS v").createOrReplaceTempView("t")
+ val query = caseWhenOverT(300, end = -1)
+ for (backend <- Seq("janino", "jdk")) {
+ withSQLConf(SQLConf.CODEGEN_COMPILER.key -> backend) {
+ // The warning is once per JVM and INFO after it, so the test reads
every level, of the
+ // logger the compilers report under.
+ val logs = new LogAppender("the trial compile's reports")
+ withLogAppender(logs,
+ loggerNames = Seq(classOf[CodeGenerator[_, _]].getName),
+ level = Some(Level.DEBUG)) {
+ assert(splitsCaseWhen(sql(query)))
+ sql(query).collect()
+ }
+
assert(!logs.loggingEvents.exists(_.getMessage.getFormattedMessage.contains(
+ "too long to be JIT compiled")), backend)
+ }
+ }
+ }
+ }
+
+ test("SPARK-33301: the code a stage keeps unsplit reports its method past
the JIT limit") {
+ // The trial compile holds its reports back for the code kept. Here the
default limit keeps
+ // the split of 300 branches, whose code in one piece was compiled in a
trial and is not kept.
+ // That code then leaves the compile cache, so a later compile of it
reports its method past
+ // the JIT limit, as for any compiled code: with the split turned off, and
with the limit
+ // raised past HotSpot's 8000 bytes, where the stage keeps the code in one
piece although the
+ // JIT leaves its method interpreted.
+ withTempView("t") {
+ spark.range(100).selectExpr("id AS v").createOrReplaceTempView("t")
+ val query = caseWhenOverT(300, end = -7)
+ def reports(confs: (String, String)*): Boolean = {
+ val logs = new LogAppender("the kept code's reports")
+ withLogAppender(logs,
+ loggerNames = Seq(classOf[CodeGenerator[_, _]].getName),
+ level = Some(Level.DEBUG)) {
+ // Compiled here rather than by running the query, which compiles
with the session's
+ // class loader, and so with Janino whatever the backend
(`CodeCompiler.active`).
+ withSQLConf(confs:
_*)(genCode(sql(query)).foreach(CodeGenerator.compile))
+ }
+ logs.loggingEvents.exists(_.getMessage.getFormattedMessage.contains(
+ "too long to be JIT compiled"))
+ }
+ for (backend <- Seq("janino", "jdk")) {
+ withSQLConf(SQLConf.CODEGEN_COMPILER.key -> backend) {
+ assert(splitsCaseWhen(sql(query)))
+ assert(reports(SQLConf.WHOLESTAGE_SPLIT_EXPRESSIONS.key -> "false"),
backend)
+ assert(reports(SQLConf.WHOLESTAGE_SPLIT_EXPRESSIONS_METHOD_LIMIT.key
-> "65535"), backend)
+ withSQLConf(SQLConf.WHOLESTAGE_SPLIT_EXPRESSIONS_METHOD_LIMIT.key ->
"65535") {
+ assert(!splitsCaseWhen(sql(query)))
+ }
+ }
+ }
+ }
+ }
+
+ test("SPARK-33301: which code a stage keeps, from its trial compiles") {
+ // `chooseCode`'s decisions, with the compile replaced, so that a method
can be as large as a
+ // decision needs and a split can fail to compile.
+ withTempView("t") {
+ spark.range(100).selectExpr("id AS v").createOrReplaceTempView("t")
+ val stage = sql(caseWhenOverT(300, end =
-1)).queryExecution.executedPlan.collectFirst {
+ case w: WholeStageCodegenExec => w
+ }.get
+ val recorded = stage.generate(WholeStageSplit.Record)
+ val ctx = recorded._1
+ assert(ctx.wholeStageSplitsRecorded.length === 1)
+ // The method the CASE WHEN's code lands in.
+ val Seq(method) =
ctx.wholeStageSplitMarkersByFunction.get(0).toSeq.flatten
+ val failure = new IllegalStateException("does not compile")
+ // Each case its own decision, which is remembered by the digest of the
marked code.
+ var copies = 0
+ class Run(
+ whole: Try[Option[Map[String, Int]]],
+ split: Try[Option[Map[String, Int]]],
+ splitIsWhole: Boolean = false) {
+ copies += 1
+ private val digest = s"copy $copies"
+ val code = (ctx, new CodeAndComment(recorded._2.body + s"\n/* $digest
*/",
+ recorded._2.comment))
+ val splitCode = if (splitIsWhole) {
+ code
+ } else {
+ (ctx, new CodeAndComment(code._2.body + "\n/* split */",
code._2.comment))
+ }
+ var splitAsked = Seq.empty[Option[Set[WholeStageSplitKey]]]
+ var compiles = 0
+ var invalidated = Seq.empty[CodeAndComment]
+ def choose(throwSplitFailure: Boolean = false): CodeAndComment = {
+ ctx.wholeStageSplitDigest = digest
+ stage.chooseCode(code,
+ only => { splitAsked :+= only; splitCode },
+ (c, _) => {
+ compiles += 1
+ WholeStageCodegenExec.TrialResult(if (c eq code._2) whole else
split, () => ())
+ },
+ c => invalidated :+= c, throwSplitFailure)._2
+ }
+ }
+ def sizes(m: (String, Int)*): Try[Option[Map[String, Int]]] =
Success(Some(m.toMap))
+ val notCompiling = Failure(failure)
+ val unknown = Success(None)
+ val recordedOnly = Some(Set(ctx.wholeStageSplitsRecorded.head))
+ // The split is kept where it lowers the bytes past the limit, and the
code in one piece
+ // leaves the compile cache.
+ val kept = new Run(sizes(method -> 20000), sizes(method -> 3000))
+ assert(kept.choose() eq kept.splitCode._2)
+ assert(kept.splitAsked === Seq(recordedOnly) && kept.invalidated ===
Seq(kept.code._2))
+ // A larger method that holds nothing marked does not stop the split of
the one that does.
+ val beside = new Run(sizes(method -> 9000, "other" -> 12000),
+ sizes(method -> 1500, "other" -> 12000))
+ assert(beside.choose() eq beside.splitCode._2)
+ // A method past the limit that holds nothing marked, or a method run
once, is left alone,
+ // with no split generated.
+ for (whole <- Seq(sizes(method -> 1000, "other" -> 20000),
+ sizes(method -> 1000, "init" -> 20000))) {
+ val left = new Run(whole, notCompiling)
+ assert(left.choose() eq left.code._2)
+ assert(left.splitAsked.isEmpty && left.compiles === 1)
+ }
+ // A split that leaves the bytes past the limit as they were is not
kept, and leaves the
+ // compile cache itself.
+ val same = new Run(sizes(method -> 20000), sizes(method -> 20000))
+ assert(same.choose() eq same.code._2)
+ assert(same.invalidated === Seq(same.splitCode._2))
+ // A split that comes out as the code in one piece is neither compiled
nor dropped: the code
+ // in one piece is kept, and its class stays cached for `doExecute`.
+ val identical = new Run(sizes(method -> 20000), notCompiling,
splitIsWhole = true)
+ assert(identical.choose() eq identical.code._2)
+ assert(identical.compiles === 1 && identical.invalidated.isEmpty)
+ // A split that fails to compile leaves the code in one piece, which
compiled, and is not
+ // tried again; under testing the failure is thrown.
+ val fails = new Run(sizes(method -> 20000), notCompiling)
+ assert(fails.choose() eq fails.code._2)
+ assert(fails.choose() eq fails.code._2)
+ assert(fails.compiles === 2)
+ val thrown = new Run(sizes(method -> 20000), notCompiling)
+ assert(intercept[IllegalStateException](thrown.choose(throwSplitFailure
= true)) eq failure)
+ // Code in one piece that fails to compile has every expression split
where that compiles,
+ // and the decision is remembered: the next generation asks for the same
split with no trial.
+ val broken = new Run(notCompiling, sizes(method -> 3000))
+ assert(broken.choose() eq broken.splitCode._2)
+ assert(broken.choose() eq broken.splitCode._2)
+ assert(broken.splitAsked === Seq(None, None) && broken.compiles === 2)
+ // Where the split fails too, the code in one piece is left to
`doExecute`, which reports
+ // its failure as without the split, also under testing; that is
remembered as well.
+ val bothFail = new Run(notCompiling, notCompiling)
+ assert(bothFail.choose(throwSplitFailure = true) eq bothFail.code._2)
+ assert(bothFail.choose(throwSplitFailure = true) eq bothFail.code._2)
+ assert(bothFail.compiles === 2)
+ // Sizes unknown, a second compile cache hit, keep the code in one piece
and decide nothing,
+ // so the next generation compiles again.
+ val unknownWhole = new Run(unknown, notCompiling)
+ assert(unknownWhole.choose() eq unknownWhole.code._2)
+ assert(unknownWhole.choose() eq unknownWhole.code._2)
+ assert(unknownWhole.compiles === 2 && unknownWhole.splitAsked.isEmpty)
+ val unknownSplit = new Run(sizes(method -> 20000), unknown)
+ assert(unknownSplit.choose() eq unknownSplit.code._2)
+ assert(unknownSplit.choose() eq unknownSplit.code._2)
+ assert(unknownSplit.compiles === 4)
+ }
+ }
+
+ test("SPARK-33301: a decision is not replayed under another
methodSplitThreshold") {
+ // Which CASE WHENs are recorded depends on the threshold, while the code
without the marks
+ // does not, so a decision made under one threshold names other
expressions under another.
+ // Each order of the two, in one session, runs the query right.
+ withTempView("t") {
+ spark.range(10).selectExpr("id").createOrReplaceTempView("t")
+ val a = (0 until 60).map(k => s"WHEN id = $k THEN id *
$k").mkString("CASE ", " ", " END")
+ val b = (0 until 300).map(k => s"WHEN id + 1 = $k THEN id * $k")
+ .mkString("CASE ", " ", " END")
+ for ((first, second) <- Seq(("1024", "40000"), ("40000", "1024"))) {
+ val query = s"SELECT $a AS a, $b AS b, $first AS threshold FROM t"
+ for (threshold <- Seq(first, second)) {
+ withSQLConf(SQLConf.CODEGEN_METHOD_SPLIT_THRESHOLD.key -> threshold)
{
+ checkAnswer(sql(query), withoutWholeStage(sql(query)))
+ }
+ }
+ }
+ }
+ }
+
+ test("SPARK-33301: a literal holding the text of a mark is left as it is") {
+ // A foldable pattern is written into the code as a Java literal. The
stage records its CASE
+ // WHEN and keeps its code in one piece, the marks stripped, and the
literal with them.
+ withTempView("t") {
+ spark.range(100).selectExpr("id AS v", "CAST(id AS STRING) AS s")
+ .createOrReplaceTempView("t")
+ val caseWhen = (0 until 64).map(k => s"WHEN v = $k THEN v * $k")
+ .mkString("CASE ", " ", " ELSE -1 END")
+ def df: DataFrame = sql(s"SELECT $caseWhen AS r, s RLIKE
'/*wholeStageSplit:0*/' AS m " +
+ "FROM t")
+ assert(recordedSplits(df) > 0 && !splitsCaseWhen(df))
+ checkAnswer(df, withoutWholeStage(df))
+ checkAnswer(sql(s"SELECT count(*) FROM t WHERE s RLIKE
'/*wholeStageSplit:0*/' AND " +
+ s"$caseWhen > -2"), Row(0))
+ }
+ val code = "a(); /*wholeStageSplit:0*/ b(\"/*wholeStageSplit:1*/\", '\\'')
// x\n" +
+ "/* other */ c(); /*wholeStageSplit:2*/"
+ assert(CodeGenerator.stripWholeStageSplitMarkers(code) ===
+ "a(); b(\"/*wholeStageSplit:1*/\", '\\'') // x\n/* other */ c(); ")
+ }
+
+ test("SPARK-33301: a simple CASE over a collated column is split") {
+ // The collated column is in every condition, so it is a common
subexpression, whose code is
+ // the column's: the split methods take the column.
+ withTempView("t") {
+ spark.range(100).selectExpr("CAST(id % 20 AS STRING) AS
s").createOrReplaceTempView("t")
+ val caseWhen = (0 until 300).map(k => s"WHEN 'k$k' THEN $k")
+ .mkString("CASE s COLLATE UNICODE_CI ", " ", " ELSE 0 END")
+ def df: DataFrame = sql(s"SELECT $caseWhen AS r FROM t")
+ withSQLConf(SQLConf.WHOLESTAGE_SPLIT_EXPRESSIONS.key -> "false") {
+ assert(maxStageMethodSize(df) >
CodeGenerator.DEFAULT_JVM_HUGE_METHOD_LIMIT)
+ }
+ assert(maxStageMethodSize(df) <=
CodeGenerator.DEFAULT_JVM_HUGE_METHOD_LIMIT)
+ checkAnswer(df, withoutWholeStage(df))
+ }
+ }
+
+ test("SPARK-33301: a filter's copy of a projection's CASE WHEN is split only
where it is past " +
+ "the limit") {
+ // The filter is pushed below the projection with its alias replaced, so
the filter evaluates
+ // a copy of the projection's CASE WHEN, in a method under the limit; the
projection's copy
+ // is in its own method, past the limit, and only that copy is split.
+ withTempView("t") {
+ spark.range(100).selectExpr("id AS v").createOrReplaceTempView("t")
+ def caseWhen(n: Int, base: Int): String =
+ (0 until n).map(i => s"WHEN v = $i THEN v * ${base +
i}").mkString("CASE ", " ", " END")
+ def query: DataFrame = sql(s"SELECT * FROM (SELECT v, ${caseWhen(64,
7000)} AS b, " +
+ s"${caseWhen(300, 9000)} AS c FROM t) WHERE b <> -1")
+ val (ctx, _) = query.queryExecution.executedPlan.collectFirst {
+ case w: WholeStageCodegenExec => w
+ }.get.doCodeGen()
+ val holding = ctx.functionsHolding("7001L")
+ assert(holding.exists(_.contains("caseWhen")), holding)
Review Comment:
**[Test] These assertions also pass with the opposite outcome, and two other
new tests don't pin what they are named for either.**
If the filter's copy were split and the projection's kept whole,
`functionsHolding("7001L")` would hold `filter_caseWhen_N_0` and
`project_doConsume_0`, and both assertions would still pass. Could this check
the operator prefix instead, e.g. that every holder containing `caseWhen`
starts with `project_`?
- "a decision is not replayed under another methodSplitThreshold" (L2133)
only checks answers: its 40000 -> 1024 order passed on the previous head even
though the wrong CASE WHEN was split, and the out-of-range check in `codeOf`
never runs. Asserting which CASE WHEN is split in the second run, e.g. that
`functionsHolding("299L")` holds `caseWhen` methods, would pin it.
- Nothing runs two threads into one undecided key, the case item 3 of my
previous review was about; the only concurrent `doCodeGen` test (SPARK-57196 in
`UnionCodegenSuite`) has no CASE WHEN. A latch in the faked compile, with the
second thread seen `WAITING`, would cover it.
- `caseWhen(n, base)` and the `collectFirst { case w: WholeStageCodegenExec
=> w }.get.doCodeGen()` lookup are repeated in two tests (L1938-1944 and
L2197-2203), and `functionsHolding` is a test-only method on `CodegenContext`
(`CodeGenerator.scala` L1526); matching function names by operator prefix in
`genCode(df)` could replace it.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/WholeStageCodegenExec.scala:
##########
@@ -610,6 +616,90 @@ object WholeStageCodegenExec {
// Reset generation time of Java source code.
// Visible for testing
def resetCodeGenTime(): Unit = _codeGenTime.set(0L)
+
+ /**
+ * A trial compile's outcome for `WholeStageCodegenExec.chooseCode`: the
bytecode size of each
+ * method, None where they are unknown, or the failure; and the reports and
metric updates the
+ * compile held back for the code kept.
+ */
+ private[execution] case class TrialResult(
+ methodSizes: Try[Option[Map[String, Int]]],
+ report: () => Unit)
+
+ /**
+ * Compiles `code` as a trial (`CodeCompiler.trial`). A hit in the compile
cache produces no
+ * method sizes, so the class is then dropped and compiled again. A hit
happens where the class
+ * is cached while the decision for its code is not: the decision was
evicted, or the same body
+ * was compiled earlier without being decided - with the split off, whose
code is the recorded
+ * code without its marks, or under another limit. Should another thread
load the class again in
+ * between, the second compile is a hit as well, and the sizes are unknown.
+ */
+ private[execution] def trialCompile(
+ code: CodeAndComment,
+ failureExpected: Boolean): TrialResult = {
+ def once(): (Try[Map[String, Int]], CodeCompiler.TrialCompile) = {
+ val (result, trial) =
CodeCompiler.trial(failureExpected)(CodeGenerator.compile(code))
+ (result.map(_ => trial.methodSizes), trial)
+ }
+ val (sizes, trial) = once() match {
+ case (Success(m), _) if m.isEmpty =>
+ CodeGenerator.invalidateCompiled(code)
+ once()
+ case result => result
+ }
+ TrialResult(sizes.map(m => Option.when(m.nonEmpty)(m)), () =>
trial.report())
+ }
+
+ /** The methods a class runs once, whose being interpreted costs nothing per
row. */
+ private val RunOnce = Set("<init>", "<clinit>", "init")
Review Comment:
**[Cleanup] `RunOnce` copies `CodeCompiler.runOnceMethods`, and
`decisionKey` builds the compile cache's key prefix by hand.**
`CodeCompiler.scala` L331 has the same set; neither covers the `init_N_M`
methods split out of `initMutableStates`, which is harmless for the decision
today, but the two copies have to stay in sync. `decisionKey` (L691-693) builds
`(CodeGenerator.classLoaderKey(), CodeCompiler.active(code), ...)` as
`CodeGenerator.cacheKey` (`CodeGenerator.scala` L2635-2637) does, so a part
added to the cache key later would be missed here silently; the fix for item 13
of my previous review shared it only within `CodeGenerator`. Could each come
from one `private[sql]` helper?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -1148,6 +1291,145 @@ class CodegenContext extends Logging {
}
}
+ /**
+ * What the code generated for `roots` reads from the scope it was generated
in, for moving that
+ * code into a method: the local variables to pass and the inputs to
evaluate first, or None
+ * where `policy` refuses the method. The code reads its inputs from
`currentVars` and, where
+ * the operator sets one, from `INPUT_ROW`; from the states in `subExprs`;
and, inside a `With`,
+ * from the definitions its references compute. A field or a literal is read
as it stands.
+ *
+ * Every place that moves generated code into a method walks the code's
expressions with this,
+ * so that a node binding a variable needs one change here; where those
places need different
+ * answers, `policy` says which (see [[InputPolicy]]).
+ */
+ private[sql] def collectInputs(
+ roots: Seq[Expression],
+ policy: InputPolicy,
+ subExprs: Map[ExpressionEquals, SubExprEliminationState]):
Option[CollectedInputs] = {
+ val args = mutable.LinkedHashMap.empty[String, VariableValue]
+ val inputsToEvaluate = mutable.ArrayBuffer.empty[ExprCode]
+ // Whether the code can read `v` from the method: false for a value no
parameter can carry
+ // where the policy asks for names. `ExpandExec` hands out a
`VariableValue` naming a slot of a
+ // compacted mutable state array, and a `SimpleExprValue` is an expression
rather than a name
+ // -- `posexplode_outer` gives its position the nullness `index == -1`,
and a `Byte` or `Short`
+ // literal's value is `(byte)1`. A field or a literal needs no parameter
and is read as it
+ // stands; so is a field handed out as a `VariableValue` where the policy
says so, passing one
+ // being refused by the split methods' own check, a slot of a compacted
array included.
+ def take(v: ExprValue): Boolean = v match {
+ case local: VariableValue =>
+ val name = local.variableName
+ val isName = name.nonEmpty &&
Character.isJavaIdentifierStart(name.head) &&
+ name.forall(Character.isJavaIdentifierPart)
+ // A slot of a compacted array is a field, since no parameter can be
named `array[i]`,
+ // where the whole stage split is on; with it off, every policy
answers as it did before
+ // that split existed. Another field is passed where the policy says
so.
+ val isField = (isCompactedSlot(name) &&
+ (policy.fieldsNeedNoArgument || wholeStageSplit !=
WholeStageSplit.Off)) ||
+ (policy.fieldsNeedNoArgument && mutableStateNames.contains(name))
+ val taken = isName || isField || !policy.onlyNames
+ if (taken && !isField) {
+ args.getOrElseUpdate(name, local)
+ }
+ taken
+ case _: GlobalValue | _: LiteralValue => true
+ case _ => !policy.onlyNames
+ }
+ var possible = !policy.rowWherever || INPUT_ROW == null ||
+ take(JavaCode.variable(INPUT_ROW, classOf[InternalRow]))
+ var readsRow = false
+ val visited = mutable.HashSet.empty[Long]
+ // The definitions of a `With` in the trees walked here, which reach
`currentCommonExprs` only
+ // once that `With` is generated. Ids come from one counter, so one map
serves every scope.
+ val nestedDefs = mutable.HashMap.empty[Long, Expression]
+ val toVisit = mutable.Stack[Expression](roots: _*)
+ while (possible && toVisit.nonEmpty) {
+ val next = toVisit.pop()
+ // A state is looked up before anything else; no input is a common
subexpression, since
+ // `EquivalentExpressions` skips leaves.
+ subExprs.get(ExpressionEquals(next)) match {
+ case Some(_) if policy.onlyStatesRead && !readsSubExprState(next) &&
Review Comment:
**[Cleanup] `!readsSubExprState(next)` is implied by
`generatedChild(next).isDefined`, and the `InputPolicy` docs still describe
refusing these nodes.**
`generatedChild` is defined only for `Alias`, `Collate` and a `Cast` with
`generatesChildCode`, and `readsSubExprState` (L1417-1420) is false for each of
them: its `Cast` arm is `!c.generatesChildCode`, and `Alias` and `Collate`
override `genCode`. So the middle conjunct never decides, and
`generatedChild(next)` is computed twice. The `onlyStatesRead` doc (L120-121,
"refusing one for a node that overrides `genCode` and may generate its child
instead") and the whole-stage policy's rationale (L172-173, "which stops at a
state too") still describe the previous behavior. Could the case compute the
child once, e.g. from `generatedChild(next).filter(_ =>
policy.onlyStatesRead)`, and the docs say that such a node is looked through?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeCompiler.scala:
##########
@@ -334,6 +335,97 @@ object CodeCompiler extends Logging {
/** Lets a test see the first warning again. */
private[catalyst] def resetHugeMethodWarning(): Unit =
hugeMethodWarned.set(false)
+ /**
+ * A compile whose result its caller may discard ([[trial]]). What the
compile would report of
+ * methods past the JIT limit is held here, for the caller to report if it
keeps the code, and so
+ * are its updates of the codegen metrics, which describe the code that runs.
+ *
+ * @param failureExpected whether a failure to compile is an answer the
caller asks for, which the
+ * compile then logs at debug level rather than as an
error.
+ */
+ private[sql] final class TrialCompile private[codegen] (val failureExpected:
Boolean) {
+ private val hugeMethods = mutable.ArrayBuffer.empty[(String, String, Int)]
+ private val methods = mutable.ArrayBuffer.empty[(String, Int)]
+ private val metricUpdates = mutable.ArrayBuffer.empty[() => Unit]
+
+ private[codegen] def holdMetricUpdate(update: () => Unit): Unit =
metricUpdates += update
+
+ private[codegen] def holdHugeMethod(className: String, methodName: String,
size: Int): Unit =
+ hugeMethods += ((className, methodName, size))
+
+ private[codegen] def recordMethod(methodName: String, size: Int): Unit =
+ methods += ((methodName, size))
+
+ /**
+ * The bytecode size of every method the compile produced, by its name in
the source, the
+ * largest where a name is in more than one class; empty where the compile
was a hit in the
+ * compile cache, which produces nothing. Janino names a private method of
an inner class with
+ * a `$` after its name in the source.
+ */
+ def methodSizes: Map[String, Int] =
+ methods.groupMapReduce(_._1.stripSuffix("$"))(_._2)(math.max)
+
+ /** Makes the reports and metric updates the compile held back, for the
code kept. */
+ def report(): Unit = {
+ metricUpdates.foreach(_())
+ hugeMethods.foreach { case (className, methodName, size) =>
+ logHugeMethod(className, methodName, size)
+ }
+ }
+ }
+
+ private val currentTrial = new ThreadLocal[TrialCompile]
+
+ /**
+ * Runs `body`, a compile whose result its caller may discard, holding back
its reports of methods
+ * past the JIT limit until the caller keeps the code
([[TrialCompile.report]]); and, where
+ * `failureExpected`, its report of a failure to compile, which is then the
answer the caller is
+ * asking for. Whole-stage codegen compiles a stage this way to decide
whether to split its
+ * expressions
(`spark.sql.codegen.wholeStage.splitExpressions.methodLimit`). The trial is the
+ * calling thread's, and a backend that compiles on a thread of its own
carries it there
+ * ([[withTrial]]), as the JDK backend does.
+ */
+ private[sql] def trial[T](failureExpected: Boolean)(body: => T): (Try[T],
TrialCompile) = {
+ val trial = new TrialCompile(failureExpected)
+ (withTrial(trial)(Try(body)), trial)
+ }
+
+ /** The trial the calling thread compiles under, or null outside one. */
+ private[codegen] def activeTrial: TrialCompile = currentTrial.get
+
+ /**
+ * Updates the codegen metrics, or, in a trial, holds the update back until
the caller keeps the
+ * code ([[TrialCompile.report]]), so that code compiled only to be
discarded is not counted.
+ */
+ private[codegen] def updateMetrics(update: => Unit): Unit =
Option(currentTrial.get) match {
Review Comment:
**[Cleanup] `TrialCompile` holds deferred reports in two channels.**
The huge-method reports are held as tuples and replayed through
`logHugeMethod`, which has its own `currentTrial` branch (L440-444), while the
metric updates are held as closures through this method. If the call in
`computeByteCodeStats` (L488) were `updateMetrics(logHugeMethod(...))`,
`holdHugeMethod`, the tuple buffer and that branch could go, and `report()`
would replay one list. That would also keep the order a compile outside a trial
has: today `report()` replays every metric update, including the loader's "Code
generated in" log, before the huge-method reports.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -1823,6 +2444,96 @@ object CodeGenerator extends Logging {
// method which is going to be called by the outer class instead of the many
small ones
final val MERGE_SPLIT_METHODS_THRESHOLD = 3
+ //
---------------------------------------------------------------------------------------------
+ // The whole stage split of expressions
(`CodegenContext.splitExpressionsWithSources`)
+ //
---------------------------------------------------------------------------------------------
+
+ /** A slot of a compacted mutable state array, such as
`mutableStateArray_0[3]`. */
+ private[codegen] val CompactedSlot = "([A-Za-z_$][\\w$]*)\\[\\d+\\]".r
+
+ /** A name and an opening parenthesis: a call, or a keyword such as `if`. */
+ private[codegen] val CallSite = "(?<![\\w$])([A-Za-z_$][\\w$]*)\\s*\\(".r
+
+ /**
+ * The comment marking the code of the expression of index `i` in
+ * `CodegenContext.wholeStageSplitsRecorded` under `WholeStageSplit.Record`.
+ */
+ private[sql] def wholeStageSplitMarker(i: Int): String =
s"$WholeStageSplitMarkerStart$i*/"
+
+ private val WholeStageSplitMarkerStart = "/*wholeStageSplit:"
+
+ private val WholeStageSplitMarker = """/\*wholeStageSplit:(\d+)\*/""".r
+
+ /** Whether `code` may hold a `wholeStageSplitMarker`, a check before a
scan. */
+ private[codegen] def holdsWholeStageSplitMarker(code: String): Boolean =
+ code.contains(WholeStageSplitMarkerStart)
+
+ /**
+ * Each `wholeStageSplitMarker` in the generated Java `code`, as its index
and where it starts and
+ * ends: a block comment that is exactly a marker. String and character
literals and other
+ * comments are skipped, so that text a literal holds, such as a user's
regular expression, is
+ * never taken for a marker.
+ */
+ private[sql] def wholeStageSplitMarkersIn(code: String): Seq[(Int, Int,
Int)] = {
Review Comment:
**[Cleanup] This is a second hand-written scanner of generated Java, next to
`JdkCodeCompiler.rewriteInnerClassRefs`.**
Both walk the source character by character through the same states (string
and char literals with backslash escapes, `//` to the end of the line, and `/*
*/`): `rewriteInnerClassRefs` in `CodeCompiler.scala` L803-864, with its own
tests in `CodeCompilerSuite`, and this one. A fix in one won't reach the other.
Could they share one scanner of spans (code, literal, comment), the marker scan
taking the block comments that are exactly a marker and the rewrite taking the
code spans? The header scanner from item 12 of my previous review is gone, but
a lexer of its own is still here.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -1823,6 +2444,96 @@ object CodeGenerator extends Logging {
// method which is going to be called by the outer class instead of the many
small ones
final val MERGE_SPLIT_METHODS_THRESHOLD = 3
+ //
---------------------------------------------------------------------------------------------
Review Comment:
**[Nit] This section still splits the constants, and nothing ends it.**
Thanks for moving the helpers into a group of their own (item 15 of my
previous review). The header now sits between `MERGE_SPLIT_METHODS_THRESHOLD`
(L2445) and `MAX_BYTES_PER_SPLIT_CALL` (L2544), `SPLIT_CALLS_PER_GROUP` (L2551)
and the class-size thresholds after them, which were one run on master; and as
the only section header in the object, it makes the remaining ~590 lines read
as part of this section. Per `AGENTS.md` ("put the new member with the code it
belongs with ... splitting a section of unrelated code in the process"), could
the group, header included, move to the end of the object, or at least below
the constants?
--
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]
