dongjoon-hyun commented on code in PR #59225:
URL: https://github.com/apache/spark/pull/59225#discussion_r4186793922
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeCompiler.scala:
##########
@@ -386,6 +467,7 @@ object CodeCompiler extends Logging {
method.getAttributes.collect { case attr: CodeAttribute =>
val byteCodeSize = attr.code.length
CodegenMetrics.METRIC_GENERATED_METHOD_BYTECODE_SIZE.update(byteCodeSize)
Review Comment:
**[Low] Discarded trial compiles still feed the codegen metrics.**
The trial holds back the huge-method report, but the compile still updates
`METRIC_GENERATED_METHOD_BYTECODE_SIZE` here, and the cache loader updates the
source-size and compile-time metrics and logs "Code generated in"
(`CodeGenerator.scala` L2591-2594). So for every stage the gate splits, the
driver's `generatedMethodSize` histogram records the unsplit method past 8000
bytes, although only the split code runs, and the compile counts double or
triple on the first execution. Could the size metrics be held in the trial like
the report, and recorded only for the code that is kept?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -1813,6 +2365,91 @@ object CodeGenerator extends Logging {
// beyond which methods will be rejected from JIT compilation
final val DEFAULT_JVM_HUGE_METHOD_LIMIT = 8000
+ /** 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:"
+
+ /** Whether `code` holds a `wholeStageSplitMarker`. */
+ private[codegen] def holdsWholeStageSplitMarker(code: String): Boolean =
+ code.contains(WholeStageSplitMarkerStart)
+
+ private val WholeStageSplitMarker = """/\*wholeStageSplit:(\d+)\*/""".r
+
+ /** `source` without the comments `wholeStageSplitMarker` adds. */
+ private[sql] def stripWholeStageSplitMarkers(source: String): String =
Review Comment:
**[Medium] The marker strip also rewrites Java string literals, so a stage
the gate keeps whole can run with user text changed.**
This is a plain regex replace over the whole stage source
(`WholeStageCodegenExec.scala` L898), while `wholeStageSplitMarkersByMethod`
deliberately skips string literals. Some user text is inlined into the source
as a Java literal: a foldable `RLIKE` pattern only goes through `escapeJava`
(`regexpExpressions.scala` L547-549), which leaves `/` and `*` as they are. For
example, with `t` = `spark.range(100).selectExpr("id AS v", "CAST(id AS STRING)
AS s")`:
```sql
SELECT CASE WHEN v = 0 THEN v * 0 ... WHEN v = 63 THEN v * 63 ELSE -1 END AS
r,
s RLIKE '/*wholeStageSplit:0*/' AS m
FROM t
```
The 64-branch CASE WHEN is recorded and the gate keeps the code whole, so
`init()` compiles `Pattern.compile("", 0)`, and `m` is true on every row
instead of false. As a `WHERE` condition, it returns all 100 rows instead of
none. The answer is right with
`spark.sql.codegen.wholeStage.splitExpressions=false`, with `methodLimit=0`,
without whole-stage codegen, or when the split is kept, so it depends on the
gate's decision. Likewise, `from_utc_timestamp(ts, 'GMT/*wholeStageSplit:1*/')`
fails on master with an invalid zone ID but silently uses GMT here
(`datetimeExpressions.scala` L2787-2789).
The marker text has to appear in a literal, so this is unlikely by accident,
but the result is silently wrong. Could the strip skip literals the way the
scanner does, or remove the markers only where they were inserted?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/WholeStageCodegenExec.scala:
##########
@@ -668,11 +735,116 @@ 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 are
+ * made for the code kept.
+ *
+ * - Where it fails to compile, the split is the only code that can run,
and every expression is
+ * 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, 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 the code's digest, so that a stage
generated again, as every
+ * execution of a query generates it, skips the trial compiles; and 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[Seq[Seq[Expression]]]] => (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)))
+ val key = decisionKey(whole, limit)
+ Option(splitDecisions.getIfPresent(key)) match {
+ case Some(SplitThese(indices)) => return splitOf(indices)
Review Comment:
**[Medium] A remembered split decision is replayed as positions into a
different recording, which can fail the query with
`IndexOutOfBoundsException`.**
`splitDecisions` is JVM-wide, and its key is the class loader, the backend,
the limit and a digest of the marker-stripped code (L677-679), while
`SplitThese` holds indices into this generation's `wholeStageSplitsRecorded`
(L795-796). Which CASE WHENs get recorded depends on
`spark.sql.codegen.methodSplitThreshold` (`CodeGenerator.scala` L1584-1585),
but a CASE WHEN that is not recorded generates the same code once the markers
are stripped, so changing the threshold changes the recording without changing
the key. For example, in one session, with `a` a CASE WHEN of 60 branches over
`id` and `b` one of 300 branches over `id + 1` (`WHEN id = k THEN id * k`):
1. `SELECT <a> AS a, <b> AS b FROM range(10)` at the default threshold
records both (`a` = 0, `b` = 1). Both are in the projection's method past 8000
bytes, so `SplitThese(Some(Set(0, 1)))` is remembered.
2. After `SET spark.sql.codegen.methodSplitThreshold=40000` (any value
between the sizes of the pieces of `a` and `b`), the same query records only
`b`, as index 0. The stripped code and the key are the same, so this line
evaluates `Set(0, 1).map(ctx.wholeStageSplitsRecorded)`, which throws
`IndexOutOfBoundsException`.
`doExecute` calls `doCodeGen()` before its compile `try` (L919-921), so the
codegen fallback doesn't apply and the query fails; `df.explain("codegen")` and
`debugCodegen()` throw as well. Because the memo is JVM-wide, two sessions with
different thresholds that share a class loader (e.g. on the Thrift server)
collide the same way. In the other direction (decided at 40000, replayed at the
default), index 0 names `a`, so the small CASE WHEN is split and the 300-branch
one stays in an interpreted method, without any trial compile.
Could the key include what the recording depends on, or could the decision
keep the chosen sources (or be checked against the current recording) instead
of positions?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -1148,6 +1273,128 @@ 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 under every policy, since no
parameter can be
+ // named `array[i]`; another field is passed where the policy says so.
+ val isField = isCompactedSlot(name) ||
Review Comment:
**[Low] `spark.sql.codegen.wholeStage.splitExpressions=false` doesn't give
back master's code.**
The compacted-slot rule here applies under every policy, so it also changes
`CommonExprSlots.methodArgs` and `getLocalInputVariableValues`, which the conf
doesn't gate. In master, `canPass` refused a slot such as
`expand_mutableStateArray_0[0]`, since it isn't a name, so a `With` definition
reading it stayed inline; now that definition becomes a `computeCommonExpr`
method that reads the slot as a field, with the conf on or off (e.g. when
`ExpandExec` hands out the value of an `UNPIVOT` over string columns as such a
slot, and a nested `nullif` in a later CASE WHEN branch of a filter right above
reads it). Likewise, operator methods no longer receive such slots as (invalid)
parameters, which is a fix. The new code looks right, but the kill switch can't
turn it off, so "With it off, the generated code is exactly that of master" and
"`getLocalInputVariableValues` keeps its signature and its results" in the
description don't hold. Could the description say so, or could the rule follow
the conf?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -1915,6 +2552,20 @@ object CodeGenerator extends Logging {
throw e.getCause
}
+ /**
+ * The class loader the compile cache keys a class on, held weakly: the
calling thread's context
+ * class loader, or Spark's.
+ */
+ private[sql] def classLoaderKey(): AnyRef =
+ new HashableWeakReference(Utils.getContextOrSparkClassLoader)
+
+ /** Drops `code`'s class from the compile cache: a variant whole stage
codegen did not keep. */
+ private[sql] def invalidateCompiled(code: CodeAndComment): Unit = {
Review Comment:
**[Cleanup] The compile-cache key is built separately here and in
`compile`.**
This spells out `(new
HashableWeakReference(Utils.getContextOrSparkClassLoader),
CodeCompiler.active(code), code)` again instead of sharing it with `compile`
(L2538-2546), and doesn't use the new `classLoaderKey()` either.
`NonFateSharingCache.invalidate` takes `Any`, so if `compile`'s key changes (it
already gained the backend once), this would still compile but silently
invalidate nothing; `trialCompile`'s retry would then get empty sizes, and
`chooseCode` would remember KeepWhole for every stage that is already cached.
Could both use one typed private key helper?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -1813,6 +2365,91 @@ object CodeGenerator extends Logging {
// beyond which methods will be rejected from JIT compilation
final val DEFAULT_JVM_HUGE_METHOD_LIMIT = 8000
+ /** 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:"
+
+ /** Whether `code` holds a `wholeStageSplitMarker`. */
+ private[codegen] def holdsWholeStageSplitMarker(code: String): Boolean =
+ code.contains(WholeStageSplitMarkerStart)
+
+ private val WholeStageSplitMarker = """/\*wholeStageSplit:(\d+)\*/""".r
+
+ /** `source` without the comments `wholeStageSplitMarker` adds. */
+ private[sql] def stripWholeStageSplitMarkers(source: String): String =
+ WholeStageSplitMarker.replaceAllIn(source, "")
+
+ /**
+ * The methods of the generated Java `source` that hold each
`wholeStageSplitMarker`, by its
+ * index. A method is a body opened, in the body of a class, after a
parameter list and any
+ * `throws` clause; a block, a lambda or an anonymous class inside a method
belongs to that
+ * method. String and character literals and comments are skipped, apart
from the markers.
+ */
+ private[sql] def wholeStageSplitMarkersByMethod(source: String): Map[Int,
Set[String]] = {
Review Comment:
**[Cleanup] Could the marker-to-method map be recorded in
`addNewFunctionToClass` instead of parsing the Java source?**
`addNewFunctionToClass` already sees every function's name and code and
checks the code for markers (L899-906), and `processNext` goes through
`addNewFunction` as well (`WholeStageCodegenExec.scala` L851). Recording
`marker index -> function name` there would give the same map, since the only
code outside added functions is run-once template code (the constructor and
`init`). That would remove this lexer, `TypeKeyword`/`MethodHeader` and
`wholeStageSplitSource`, together with the risk that a future template shape
(e.g. `if (x == Foo.class) {` inside a method, which this scanner takes as a
class body) silently attributes a marker to the wrong method, so that a method
past the limit is never split. As it stands, the scan also runs on every
decision miss before checking whether `over` is empty
(`WholeStageCodegenExec.scala` L809-811). Minor: since the last commit, the
scaladoc "A method's header up to its body..." (L2447) documents `TypeKeyword`
rather than `MethodHeader`.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -1148,6 +1273,128 @@ 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 under every policy, since no
parameter can be
+ // named `array[i]`; another field is passed where the policy says so.
+ val isField = isCompactedSlot(name) ||
+ (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(state) =>
+ possible = !policy.refuseSubExprStates &&
+ (!policy.onlyStatesRead || readsSubExprState(next)) &&
+ take(state.eval.value) && take(state.eval.isNull)
+ case None => next match {
+ case ref: BoundReference if currentVars != null &&
currentVars(ref.ordinal) != null =>
+ val input = currentVars(ref.ordinal)
+ if (input.code != EmptyBlock) {
+ if (policy.hoistDeferredInputs) {
+ inputsToEvaluate += input.copy()
+ input.code = EmptyBlock
+ } else {
+ possible = false
+ }
+ }
+ possible &&= take(input.value) && take(input.isNull)
+ case _: BoundReference =>
+ readsRow = true
+ case w: With if policy.followCommonExprs =>
+ // Only `child` is generated: a definition is generated where a
reference reaches it,
+ // so one no reference reaches is not in the code, and what it
reads decides nothing.
+ w.defs.foreach(d => nestedDefs.put(d.id.id, d.child))
+ toVisit.push(w.child)
+ case ref: CommonExpressionRef if policy.followCommonExprs =>
+ // The definition this reference fills, where the reference is, so
what that
+ // definition reads is read here as well. One belonging to a
`With` in the trees walked
+ // here is in `nestedDefs`; a sibling of the definition walked, or
one of an enclosing
+ // scope, is registered in `currentCommonExprs`. Whether the
reference calls a method
+ // that takes `INPUT_ROW` is told from the code
(`wholeStageSplitArguments`).
+ if (visited.add(ref.id.id)) {
+ nestedDefs.get(ref.id.id)
+ .orElse(currentCommonExprs.get(ref.id.id).map(_.definition))
+ .foreach(toVisit.push)
+ }
+ case e => toVisit.pushAll(e.children)
+ }
+ }
+ }
+ val arguments = args.values.toSeq
+ Option.when(possible && (!policy.checkParamLength ||
+ isValidParamLength(calculateParamLengthFromExprValues(arguments)))) {
+ CollectedInputs(arguments, readsRow, inputsToEvaluate.toSeq)
+ }
+ }
+
+ /**
+ * Whether `name` is a slot of a compacted mutable state array, such as
`mutableStateArray_0[3]`,
+ * which `addMutableState` hands out for a state it does not inline: a field
of the class.
+ */
+ private def isCompactedSlot(name: String): Boolean = name match {
+ case CodeGenerator.CompactedSlot(array) =>
mutableStateNames.contains(array)
+ case _ => false
+ }
+
+ /**
+ * Whether the code generated for `e` reads its subexpression elimination
state, which
+ * `Expression.genCode` does; a node that overrides `genCode` may generate
its child instead, as
+ * `Alias`, `Collate` and a `Cast` that changes nothing do, and then reads
what is below it. A
+ * `Cast` says which it does; any other override is refused, being told from
the class
+ * (`CodeGenerator.leavesGenCodeToExpression`) rather than asked, since a
method on `Expression`
+ * saying the same would default to reading the state, and so accept an
expression from outside
+ * Spark whose override generates its child.
+ */
+ private def readsSubExprState(e: Expression): Boolean = e match {
+ case c: Cast => !c.generatesChildCode
+ case _ => CodeGenerator.leavesGenCodeToExpression.get(e.getClass)
Review Comment:
**[Medium] A CASE WHEN whose conditions share a `Collate` is never split,
but it still pays the gate.**
Under `InputPolicy.wholeStageSplit`, a subexpression state whose class
overrides `genCode` refuses the whole block (L1331-1333). `Collate`, `Alias`
and an identity `Cast` generate exactly what their child reads, though. A
simple CASE puts the same operand into every condition:
```sql
SELECT CASE s COLLATE UNICODE_CI WHEN 'k1' THEN 1 WHEN 'k2' THEN 2 ... ELSE
0 END FROM t
```
`Collate(s)` is in every `EqualTo`, so subexpression elimination gives it a
state (`CaseWhen.branchGroups`), while `Collate.genCode` is just
`child.genCode(ctx)` (`collationExpressions.scala` L85) and never reads that
state. Every block is refused, so with hundreds of WHENs the method stays past
8000 bytes and runs interpreted, or the stage falls back past 64KB, i.e.
SPARK-33301 stays unfixed for this shape. Since the `ELSE` piece alone is
collectable, the CASE WHEN is still recorded, so the first execution also pays
the trial compiles, and the three compiles from my comment at
`WholeStageCodegenExec.scala` L825.
Could the walk descend into the child for these pass-through nodes?
Otherwise, could the description list this among the exceptions?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/WholeStageCodegenExec.scala:
##########
@@ -668,11 +735,116 @@ 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 are
+ * made for the code kept.
+ *
+ * - Where it fails to compile, the split is the only code that can run,
and every expression is
+ * 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, 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 the code's digest, so that a stage
generated again, as every
+ * execution of a query generates it, skips the trial compiles; and 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[Seq[Seq[Expression]]]] => (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)))
+ val key = decisionKey(whole, limit)
+ Option(splitDecisions.getIfPresent(key)) match {
+ case Some(SplitThese(indices)) => return splitOf(indices)
+ case Some(_) => return recorded
+ case None =>
+ }
+ val wholeTrial = compile(whole, true)
+ val (decision, code) = wholeTrial.methodSizes match {
+ case Failure(_) =>
+ // `doExecute` compiles the split, reporting as any compile does.
+ (SplitThese(None), splitOf(None))
+ case Success(sizes) =>
+ val over = sizes.filter { case (m, size) => size > limit &&
!RunOnce(m) }
+ val marked =
CodeGenerator.wholeStageSplitMarkersByMethod(ctx.wholeStageSplitSource)
+ val indices = marked.collect { case (i, methods) if
methods.exists(over.contains) => i }
+ if (indices.isEmpty) {
+ wholeTrial.report()
+ (KeepWhole, recorded)
+ } else {
+ val splitCode = splitOf(Some(indices.toSet))
+ val splitTrial = compile(splitCode._2, false)
+ splitTrial.methodSizes match {
+ case Success(splitSizes) if bytesOver(splitSizes, limit) <
over.values.sum =>
+ splitTrial.report()
+ invalidate(whole)
+ (SplitThese(Some(indices.toSet)), splitCode)
+ case Success(_) =>
+ wholeTrial.report()
+ invalidate(splitCode._2)
Review Comment:
**[Performance] When the split changes nothing, the stage is compiled three
times and the kept class is evicted.**
Recording accepts a CASE WHEN when any single piece is collectable
(`pieces.exists` in `wholeStageSplitPossible`), but the split refuses a block
as a whole. When every block holds a piece that reads an input not evaluated
yet, `splitWholeStage` returns the pieces unchanged (`CodeGenerator.scala`
L1524-1525), and `splitCode._2` is byte-identical to `whole`. For example, over
a Parquet table with `k, c0, ..., c90`: `CASE WHEN k = 0 THEN 0`, followed by
about 270 branches in which every block of about 1KB has a branch like `WHEN k
= i THEN c<i>` reading a column used only once, and `ELSE c0`. Then:
1. The split trial (L817) hits `whole`'s cache entry (`CodeAndComment`
equality is by body), so `trialCompile` invalidates and recompiles it.
2. The sizes are equal, so this branch runs `wholeTrial.report()` and
`invalidate(splitCode._2)`, which evicts the class being kept.
3. `doExecute` compiles it a third time, outside a trial, and logs the huge
method again.
That makes three driver compiles of a stage past 8000 bytes, and two
reports, on the first execution, while the scaladoc of
`wholeStageSplitPossible` says that erring towards yes "costs the stage a trial
compile and nothing else". Could `chooseCode` keep the whole code without the
second trial and without the invalidation when `splitCode._2 == whole`?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/WholeStageCodegenExec.scala:
##########
@@ -668,11 +735,116 @@ 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 are
+ * made for the code kept.
+ *
+ * - Where it fails to compile, the split is the only code that can run,
and every expression is
+ * 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, 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 the code's digest, so that a stage
generated again, as every
+ * execution of a query generates it, skips the trial compiles; and 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[Seq[Seq[Expression]]]] => (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)))
+ val key = decisionKey(whole, limit)
+ Option(splitDecisions.getIfPresent(key)) match {
+ case Some(SplitThese(indices)) => return splitOf(indices)
+ case Some(_) => return recorded
+ case None =>
+ }
+ val wholeTrial = compile(whole, true)
Review Comment:
**[Low] Any failure of the unsplit trial is remembered as "split
everything", and its error is only logged at DEBUG.**
The trial of `whole` runs with `failureExpected = true` for every kind of
failure, not only the size failures a split can fix:
`CodeCompiler.logCompileFailure` logs the error at DEBUG,
`logGeneratedCodeOnFailure` drops the source, and `SplitThese(None)` is
remembered (L839) and never revised when the split fails as well. For example,
take a projection with a recorded 64-branch CASE WHEN plus a `coalesce(v, v +
1, ..., v + 4000)` or a non-literal `IN` list (which still aren't split under
whole-stage codegen) that pushes the method past 64KB. The split variant fails
the same way in `doExecute` and the stage falls back, but the ERROR log shows
only the split variant, not the code that failed first. Every later execution
also pays the extra generation before the same failing compile, and under
testing a compile error of the unsplit code that the split happens to avoid is
no longer thrown. Could the original failure be reported, and the outcome
remembered, when the split fails too?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -1188,6 +1435,184 @@ class CodegenContext extends Logging {
}
}
+ /**
+ * What whole stage codegen does with an expression it could split, which
+ * `WholeStageCodegenExec.doCodeGen` sets for each generation of a stage.
Read here rather than
+ * the conf, so that every split of a generation follows the decision
`doCodeGen` made from the
+ * stage's own conf.
+ */
+ private[sql] var wholeStageSplit: WholeStageSplit = WholeStageSplit.Off
+
+ /**
+ * Under `WholeStageSplit.Record`, the expressions
[[splitExpressionsWithSources]] might have
+ * split, by the index their code is marked with: the sources of their
pieces.
+ */
+ private[sql] val wholeStageSplitsRecorded =
mutable.ArrayBuffer.empty[Seq[Seq[Expression]]]
+
+ /**
+ * Under `WholeStageSplit.Record`, the stage's source with the marks of the
recorded expressions
+ * (`CodeGenerator.wholeStageSplitMarkersByMethod`), which
`WholeStageCodegenExec.generate` sets
+ * before it strips them.
+ */
+ private[sql] var wholeStageSplitSource: String = ""
+
+ /**
+ * `splitExpressionsWithCurrentInputs` for pieces of code whose expressions
are known, which it
+ * splits under whole stage codegen as well. There the inputs are local
variables of the
+ * operator's method, and which of them a piece of code reads is known only
from the expressions
+ * it was generated for, which each piece carries.
+ *
+ * The pieces are packed into blocks as `splitExpressions` packs them, and
each block becomes a
+ * method taking the locals its own pieces read (`collectInputs`, under
+ * `InputPolicy.wholeStageSplit`). A block whose code reads a local no
method could take - an
+ * input the operator has not evaluated yet, whose code names locals of the
operator producing
+ * the row, or a value no parameter can name - stays where it was, between
the calls; so a piece
+ * must be valid both in the caller and inside a split function after
`makeSplitFunction`, as
+ * `CaseWhen`'s branches are. The calls of each run of consecutive split
blocks are grouped until
+ * at most `SPLIT_CALLS_PER_GROUP` are left in the caller: under whole stage
codegen every
+ * expression of an operator is generated into one method, so the budget per
call of
+ * `splitExpressions` would add up there (see `groupSplitCalls`).
+ *
+ * Outside whole stage codegen this is `splitExpressionsWithCurrentInputs`;
inside it, the code
+ * is split as `wholeStageSplit` says: in one piece where it is off, in one
piece and marked where
+ * it records, and split where it splits this expression.
+ *
+ * @param pieces the codes to evaluate expressions, each with the
expressions it was generated
+ * from.
+ * @param funcName the split function name base.
+ * @param extraArguments the list of (type, name) of the arguments of the
split functions,
+ * besides the current inputs.
+ * @param returnType the return type of the split function.
+ * @param makeSplitFunction makes split function body, e.g. add preparation
or cleanup.
+ * @param foldFunctions folds the split function calls.
+ */
+ private[sql] def splitExpressionsWithSources(
+ pieces: Seq[(String, Seq[Expression])],
+ funcName: String,
+ extraArguments: Seq[(String, String)] = Nil,
+ returnType: String = "void",
+ makeSplitFunction: String => String = identity,
+ foldFunctions: Seq[String] => String = _.mkString("", ";\n", ";")):
String = {
+ val codes = pieces.map(_._1)
+ if (currentVars == null) {
+ splitExpressionsWithCurrentInputs(
+ codes, funcName, extraArguments, returnType, makeSplitFunction,
foldFunctions)
+ } else {
+ wholeStageSplit match {
+ case WholeStageSplit.Record if wholeStageSplitPossible(pieces,
extraArguments) =>
+ val marker =
CodeGenerator.wholeStageSplitMarker(wholeStageSplitsRecorded.length)
+ wholeStageSplitsRecorded += pieces.map(_._2)
+ marker + codes.mkString("\n")
+ case WholeStageSplit.Split(only) if
only.forall(_.contains(pieces.map(_._2))) =>
Review Comment:
**[Low] `Split(only)` matches CASE WHENs by the equality of their sources,
not by occurrence.**
The decision is made per method (`WholeStageCodegenExec.scala` L811), but it
is applied to every CASE WHEN whose bound sources equal a chosen one, in a
separate generation. That goes wrong both ways:
- Over-match: in `SELECT * FROM (SELECT v, <64-branch CASE WHEN> AS b,
<300-branch CASE WHEN> AS c FROM t) WHERE b <> -1` (with non-literal branch
values), the filter is pushed below the project with the alias substituted, so
`FilterExec` evaluates an equal CASE WHEN inline in `processNext`. Only the
project's copy is marked in the method past the limit, but the filter's copy is
split too, although "an expression in a method the JIT compiles whole is left
alone" (the description measures that cost at 5-15% at 64 branches).
- Under-match: `generateExpressions` fresh-copies stateful expressions in
every generation, and a `NamedLambdaVariable` copy gets a new
`AtomicReference`, which takes part in case-class equality. A large CASE WHEN
holding a lambda in a hash join's stream-side key (written by the user, or
added by `NormalizeFloatingNumbers` for `array<double>` or `map<_, double>`
keys) then never matches the recorded sources, and is never split.
Could the choice be made by occurrence instead, e.g. by the order of the
calls within the generation?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -1188,6 +1435,184 @@ class CodegenContext extends Logging {
}
}
+ /**
+ * What whole stage codegen does with an expression it could split, which
+ * `WholeStageCodegenExec.doCodeGen` sets for each generation of a stage.
Read here rather than
+ * the conf, so that every split of a generation follows the decision
`doCodeGen` made from the
+ * stage's own conf.
+ */
+ private[sql] var wholeStageSplit: WholeStageSplit = WholeStageSplit.Off
+
+ /**
+ * Under `WholeStageSplit.Record`, the expressions
[[splitExpressionsWithSources]] might have
+ * split, by the index their code is marked with: the sources of their
pieces.
+ */
+ private[sql] val wholeStageSplitsRecorded =
mutable.ArrayBuffer.empty[Seq[Seq[Expression]]]
+
+ /**
+ * Under `WholeStageSplit.Record`, the stage's source with the marks of the
recorded expressions
+ * (`CodeGenerator.wholeStageSplitMarkersByMethod`), which
`WholeStageCodegenExec.generate` sets
+ * before it strips them.
+ */
+ private[sql] var wholeStageSplitSource: String = ""
+
+ /**
+ * `splitExpressionsWithCurrentInputs` for pieces of code whose expressions
are known, which it
+ * splits under whole stage codegen as well. There the inputs are local
variables of the
+ * operator's method, and which of them a piece of code reads is known only
from the expressions
+ * it was generated for, which each piece carries.
+ *
+ * The pieces are packed into blocks as `splitExpressions` packs them, and
each block becomes a
+ * method taking the locals its own pieces read (`collectInputs`, under
+ * `InputPolicy.wholeStageSplit`). A block whose code reads a local no
method could take - an
+ * input the operator has not evaluated yet, whose code names locals of the
operator producing
+ * the row, or a value no parameter can name - stays where it was, between
the calls; so a piece
+ * must be valid both in the caller and inside a split function after
`makeSplitFunction`, as
+ * `CaseWhen`'s branches are. The calls of each run of consecutive split
blocks are grouped until
+ * at most `SPLIT_CALLS_PER_GROUP` are left in the caller: under whole stage
codegen every
+ * expression of an operator is generated into one method, so the budget per
call of
+ * `splitExpressions` would add up there (see `groupSplitCalls`).
+ *
+ * Outside whole stage codegen this is `splitExpressionsWithCurrentInputs`;
inside it, the code
+ * is split as `wholeStageSplit` says: in one piece where it is off, in one
piece and marked where
+ * it records, and split where it splits this expression.
+ *
+ * @param pieces the codes to evaluate expressions, each with the
expressions it was generated
+ * from.
+ * @param funcName the split function name base.
+ * @param extraArguments the list of (type, name) of the arguments of the
split functions,
+ * besides the current inputs.
+ * @param returnType the return type of the split function.
+ * @param makeSplitFunction makes split function body, e.g. add preparation
or cleanup.
+ * @param foldFunctions folds the split function calls.
+ */
+ private[sql] def splitExpressionsWithSources(
+ pieces: Seq[(String, Seq[Expression])],
+ funcName: String,
+ extraArguments: Seq[(String, String)] = Nil,
Review Comment:
**[Cleanup] `extraArguments` is always `Nil`.**
It is threaded through `splitWholeStage`, `wholeStageBlocks`,
`wholeStageSplitPossible` and `wholeStageSplitArguments`, but `CaseWhen`, the
only caller, and the tests never pass it. A caller that followed the
`splitExpressionsWithCurrentInputs` idiom would also trip the audit:
`hash.scala` L428, for instance, passes its local `ev.value`, which
`splitCall(..., args ++ extraArguments, ..., audited = true)` would record as
an audited argument, so `assertSplitCallsWithin` would fail under testing once
an operator moves that code into a method of its own. Could the parameter wait
until a caller needs it?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/WholeStageCodegenSuite.scala:
##########
@@ -1799,4 +1803,678 @@ 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 Seq(code) = genCode(query())
+ def methodsWith(literal: String): Set[String] =
CodeGenerator.wholeStageSplitMarkersByMethod(
+ code.body.replace(literal,
CodeGenerator.wholeStageSplitMarker(0))).getOrElse(0, Set.empty)
+ 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(", ", ", ")")
+ val df = sql(s"SELECT hash($struct) + 1 AS a, hash($struct) + 2 AS b
FROM range(10)")
+ checkAnswer(df, withoutWholeStage(df))
Review Comment:
**[Test] This compares the plan without whole-stage codegen with itself.**
`df` is a `val`, and `withoutWholeStage(df)` is evaluated first: it builds
`df.queryExecution.executedPlan` with whole-stage codegen off, and since
`QueryExecution` builds it only once, `checkAnswer(df, ...)` then collects that
same plan. Only `materializedRdd.count()` inside `checkAnswer` runs a fresh
plan with whole-stage codegen, and it only counts the rows, so a wrong hash
from the rewritten `groupSplitCalls` path would still pass. The other new tests
use `def df`, a `query()` thunk or a fresh `sql(query)`; could this one be a
`def` as well?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/WholeStageCodegenExec.scala:
##########
@@ -610,6 +614,69 @@ 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, or the failure, and the reports of methods past the JIT limit
held back for the code
+ * kept.
+ */
+ private[execution] case class TrialResult(methodSizes: Try[Map[String,
Int]], report: () => Unit)
+
+ /**
+ * Compiles `code` as a trial (`CodeCompiler.trial`). A hit in the compile
cache produces no
+ * method sizes, so the code is then compiled again, which happens only
where the decision for it
+ * is no longer remembered while its class still is.
+ */
+ 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 {
Review Comment:
**[Medium] Concurrent first executions of a stage recompile it once per
thread, and the unchecked retry can decide from an empty size map.**
A compile records method sizes only when the cache loader runs on the
calling thread (`CodeCompiler.scala` L470), and `NonFateSharingCache.get` makes
a waiting thread take the class that another thread loaded
(`NonFateSharingCache.scala` L107-108). So when several threads run the same
stage before the first `splitDecisions.put` (L839), e.g. one dashboard query
from several Thrift sessions that share the class loader, each waiter gets a
hit with no sizes, invalidates the class that another thread's `doExecute` is
about to use, and compiles it again under the key lock. A stage with a method
past 8000 bytes is then compiled up to once per thread, one after another, and
its split variant likewise, where master compiles it once.
With three or more threads, a waiter's retry can hit again after another
thread has reloaded the class. The second result isn't checked, so
`Success(Map.empty)` makes `over` empty and that execution keeps the unsplit
code (L812-814), running its method past 8000 bytes interpreted; on the split
side, `bytesOver(Map.empty, limit) = 0` keeps the split without the check. A
KeepWhole put can also be the last write for the key. The scaladoc above also
misses the sequential cases, e.g. the same body compiled earlier with
`spark.sql.codegen.wholeStage.splitExpressions=false` (the stripped Record code
is the Off code) or under another `methodLimit`.
Could concurrent planners wait for one decision per key (e.g.
`NonFateSharingCache.get(key, loader)` for `splitDecisions`), and an empty
retry be treated as unknown rather than as "nothing past the limit"? Keeping
the per-method sizes with `ByteCodeStats` in the compile cache would also let a
hit answer without the invalidate and recompile.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -1813,6 +2365,91 @@ object CodeGenerator extends Logging {
// beyond which methods will be rejected from JIT compilation
final val DEFAULT_JVM_HUGE_METHOD_LIMIT = 8000
+ /** A slot of a compacted mutable state array, such as
`mutableStateArray_0[3]`. */
Review Comment:
**[Nit] These helpers split the JVM-limit constants.**
`DEFAULT_JVM_HUGE_METHOD_LIMIT`, `MAX_JVM_METHOD_PARAMS_LENGTH` and
`MAX_JVM_CONSTANT_POOL_SIZE` were adjacent, and now about 85 lines of marker
and parsing helpers sit between the first and the other two. Per `AGENTS.md`
("put the new member with the code it belongs with ... splitting a section of
unrelated code in the process"), could they move below the constants as a group
of their own, like `classLoaderKey` and `invalidateCompiled` next to the
compile cache?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/WholeStageCodegenExec.scala:
##########
@@ -668,11 +735,116 @@ 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)),
Review Comment:
**[Performance] A split stage generates its code twice on every execution,
and one with a per-run literal never hits the memo.**
`recorded` is computed before the memo lookup, so even when `SplitThese` is
remembered, every `doCodeGen` runs a full `generate(Record)` and a SHA-256 of
the code only to build the key (L797), and then `generate(Split(only))` (L799).
The description mentions this, but the committed results put it above the
quoted "4-7%" for one 300-branch CASE WHEN: "split as needed again" against
"split always" (`methodLimit=0`, one generation, the same code) is 247/232,
239/216 and 161/143 ms on JDK 17/21/25 (6.5%, 10.6% and 12.6%), and the keyed
aggregate, which the description doesn't mention, is 263/243, 265/234 and
191/167 ms.
More importantly, numeric and timestamp literals are inlined into the Java
source (`literals.scala` L604-631). A split stage with `current_timestamp()` or
`now()`, a streaming batch timestamp, or per-run numeric parameters therefore
gets a new digest on every execution. It never hits the memo, and pays the
Record generation, the whole trial compile, the Split generation and the split
trial compile each time: two generations and two driver compiles, where master
and `methodLimit=0` do one of each. A plan-keyed memo wouldn't help there
either unless literals are normalized. Could the decision be keyed on something
that doesn't change with literal values?
--
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]
