MaxGekk commented on code in PR #59225:
URL: https://github.com/apache/spark/pull/59225#discussion_r4188177008
##########
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:
Thanks, this was right. The decision is now keyed on a digest of the marked
source, taken before the markers are stripped, so the marker positions fix what
each index means. `methodSplitThreshold` is in the key as well, since the split
packs its blocks by it. A remembered `SplitThese` whose indices go past the
current recording is treated as a miss and dropped from the memo. The new test
"a decision is not replayed under another methodSplitThreshold" runs the same
query at two thresholds in one session, in both orders.
##########
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:
Thanks, you're right. The regex is replaced by
`CodeGenerator.wholeStageSplitMarkersIn`, a scan that skips string and char
literals and line comments, and takes a block comment as a marker only when it
is exactly a marker. `stripWholeStageSplitMarkers` removes only the spans it
finds. The marker-to-method map uses the same scan (your comment at L2398), so
the two cannot drift apart. The new test "a literal holding the text of a mark
is left as it is" uses your `RLIKE` example: it checks the answer with the code
kept whole, and it unit-tests the strip over a string literal, a char literal
and comments.
##########
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:
Thanks, this was right on both counts.
1. `splitDecisions` is now a `NonFateSharingCache`, and the decision is made
inside `get(key, loader)`, so threads that meet the stage at once wait for one
decision instead of each invalidating and recompiling it.
2. A trial whose sizes come back empty after the retry yields `None`
(unknown). With unknown sizes the code in one piece is kept with its report,
and nothing is remembered (`Undecided` is invalidated), on the unsplit side and
the split side alike.
The `trialCompile` scaladoc now lists the sequential hits. Keeping the
per-method sizes with `ByteCodeStats` in the compile cache would remove the
recompile on a hit altogether. But the held reports would then have to be
rebuilt from the cache, so I listed that as a follow-up.
##########
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:
Thanks, this was right. Under the whole-stage policy, `collectInputs` now
descends into the child of a node whose code is its child's (`Alias`,
`Collate`, and a `Cast` with `generatesChildCode`) instead of refusing the
node's state. That is exactly what their code reads, so an enclosing method
that compiles today already has those variables in scope. A node whose state is
a real subexpression result (e.g. `Upper`) is still taken as the state. New
tests: "a simple CASE over a collated column is split", a 300-branch simple
CASE over `s COLLATE UNICODE_CI` that checks every method stays under the
limit, and a `CollectInputsSuite` case for both kinds of node.
##########
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:
Thank you, right on both points.
**The numbers.** Against "split always", one 300-branch CASE WHEN costs 6.5
/ 10.6 / 12.6% on the repeat (7.3 / 5.1 / 13.3% on the first run), not 4-7%.
The keyed aggregate costs 8-15%, which the description had left out. The
description now says 5-13%, 5-17% and 8-15%.
**The per-run literals.** `ComputeCurrentTime` turns `current_timestamp()`
into a literal, and `Literal.doGenCode` writes it into the source. So the
digest changes on every run, and a split stage pays two generations and two
compiles each time, where master pays one of each. The description now says
this instead of claiming nothing new is compiled after the first execution.
A key that ignores literal values would avoid it: either the code digested
with numeric literals masked, or the plan with literals normalized plus every
conf that affects codegen. But literal values move bytecode sizes by a byte or
two near the limit, and a plan key has to carry the codegen confs. So I'd like
to do it in a follow-up under its own JIRA; it is in the follow-up list. Does
that work for you?
##########
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:
Good catch, thank you. `chooseCode` now compares the split code with the
code in one piece before the split trial. When they are equal, it keeps the
code in one piece with no second trial and no invalidation, so the stage is
compiled once and `doExecute` takes it from the cache. The faked-compile
`chooseCode` test covers it: one compile, nothing invalidated.
##########
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:
Thanks, both cases were real. I didn't key on the global order of calls as
it stands. `subexpressionEliminationForWholeStageCodegen` regenerates the
common subexpressions in functions only when their code is long (L2110), and a
split changes that length. So the two generations can make a different number
of calls, and every later index would shift.
Each call is now keyed instead by a `WholeStageSplitKey`: its sources with
the lambda variables' values ignored (`withoutLambdaValues`), plus its
occurrence among calls with the same sources. The occurrence counters are
rolled back where the non-split subexpression pass is discarded, so the
occurrences in the code that is kept line up in both generations. New tests:
- "a filter's copy of a projection's CASE WHEN is split only where it is
past the limit": the pushed-down filter's copy stays whole;
- "a CASE WHEN in a join key that holds lambda variables is split": an
`array<double>` broadcast join key that `NormalizeFloatingNumbers` rewrites.
--
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]
