MaxGekk commented on code in PR #59225:
URL: https://github.com/apache/spark/pull/59225#discussion_r4201482979
##########
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:
Right, thank you; I had missed that the installed core jar has Guava
relocated. `splitDecisions` is a plain Guava `Cache` again, kept within
sql/core, and `decided` serializes the deciders with core's Guava-free
`KeyLock`, as you suggested: under the key's lock it returns the remembered
decision, or decides and remembers it unless it is `Undecided`. That also
removes the catch of Guava's wrapper exceptions, since nothing wraps them any
more. The scaladoc names SPARK-44064 and SPARK-50443 so the next person doesn't
reach for `NonFateSharingCache.apply(Cache)` again.
##########
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:
Confirmed: `NamedLambdaVariable` declares no pattern, so the pruned
transform never ran. It is now a plain `transform`. The test now builds the
stream side's key from `transform(a, x -> x + k)` in each of 60 branches,
asserts that the unsplit method is between 8000 and 65535 bytes, and then that
the CASE WHEN is split. With the old pruned transform restored, the test fails
with `splitsCaseWhen(df) was false`. The description no longer mentions the
`array<double>` key.
##########
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:
Done, as you suggested. Where the code in one piece fails, the split of
every expression is compiled only if it turns some block into a method. Its
compile expects its failure only outside testing (`failureExpected =
!throwSplitFailure`). Under testing its failure is thrown, and otherwise it is
warned about with the exception, as in the other branch. Both warnings now
carry the exception. `SplitFailed` is gone, since it was handled exactly like
`KeepWhole`, and the `KeepWhole` scaladoc says what it covers. The
faked-compile test checks the warning, the `failureExpected` flags, the throw
under testing, and one compile where no block can be 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 =>
+ 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:
Agreed, comparing the code was the wrong signal.
`CodegenContext.wholeStageBlocksSplit` counts the blocks a `Split` generation
turns into methods, and subexpression elimination restores it with the
occurrences when it discards its first pass. `chooseCode` now keeps the code in
one piece, without a second trial, when the count is zero. A new test builds
your shape with real generations and compiles: a recorded CASE WHEN past the
limit, every block holding a branch that reads a deferred column, beside common
subexpressions whose discarded pass adds `elementToString` helpers. It asserts
that the split generation differs from the code in one piece, that it split no
block, and that only one compile ran. With the old `splitCode._2 == whole`
check restored, it fails with 2 compiles. The faked-compile case now uses a
split whose code differs but which splits nothing.
##########
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:
Thanks, that's a better account than mine. The `doCodeGen` scaladoc and the
conf doc now say that, with the conf on, every stage reads a compacted slot as
a field. They also say that only a generation that splits removes the discarded
pass's functions. I fixed the comment in
`subexpressionEliminationForWholeStageCodegen` the same way. The redundant
`policy.fieldsNeedNoArgument ||` is gone. A new test pins the UNPIVOT case with
your `f`: with the conf on the stage compiles and answers right, and with it
off compiling the stage fails with "')' expected instead of '['". The
description mentions the fix, corrects "removed", and drops the stale CI line.
##########
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:
Right, the data now has `concat('K', ...)`, and the test asserts that every
row takes a branch. The identity `Cast` has its own test now. `CASE CAST(s AS
STRUCT<b: BIGINT>) WHEN named_struct('b', kL) ...` over a parquet column `s` of
type `struct<a: bigint>` is past the limit unsplit, splits, and every row
matches a branch.
##########
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:
- The filter-copy test now asserts that every method holding the 64-branch
CASE WHEN's literal and named `caseWhen` starts with `project_`, and that some
other method holds it.
- The threshold test asserts, under 1024 in both orders, that the 300-branch
CASE WHEN's literal is in `caseWhen` methods.
- The faked-compile test now runs two threads into one undecided key. The
first is held in its trial compile by a latch. The second is seen `WAITING` on
the key's lock before the latch opens. Afterwards both return the split, with
two compiles in total.
- `functionsHolding` is gone from `CodegenContext`. The suite has a
`methodsHolding(code, text)` helper over the generated source, plus
`caseWhenOverV` and `stageOf`, which replace the repeated `caseWhen(n, base)`
and the stage lookups.
##########
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:
Agreed. `chooseCode` now ends as you sketched: `decided(key, ...)`, then the
decider's own code, or the decision's. `codeOf`, the fast path, the
out-of-range miss and the non-local `return` are gone. `Undecided` is never
stored, so it needs no invalidation.
##########
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:
Done. The case computes `generatedChild(next).filter(_ =>
policy.onlyStatesRead)` once, and the `onlyStatesRead` and whole-stage policy
docs now say that such a node is looked through.
##########
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:
Done. `CodeCompiler.forEachJavaSpan` is now the one scanner. It calls back
with each span of code, literal or comment, and the spans cover the source.
`rewriteInnerClassRefs` rewrites the code spans and copies the rest, and
`wholeStageSplitMarkersIn` reads the comments. `rewriteInnerClassRefs`'
existing tests pass unchanged, and `CodeCompilerSuite` has a test of the spans,
unterminated ones included.
##########
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:
Done. `computeByteCodeStats` calls `reportOrHold(logHugeMethod(...))`, the
renamed `updateMetrics`. `holdHugeMethod`, the tuple buffer and the trial
branch in `logHugeMethod` are gone. `report()` replays one list in the order
made, so a kept trial reports as a compile outside a trial does.
--
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]
