MaxGekk commented on code in PR #59225:
URL: https://github.com/apache/spark/pull/59225#discussion_r4188178786


##########
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:
   Thanks, agreed on the main point. When the unsplit trial fails, the split is 
now trial-compiled too:
   - if the split compiles, it is kept (`SplitThese(None)`), and `doExecute` 
finds it in the cache;
   - if it fails as well, the code in one piece is remembered (`KeepWhole`), so 
`doExecute` compiles it outside a trial. It then logs, falls back or throws its 
original error exactly as master does, and later executions generate the stage 
once.
   
   The faked-compile test covers both branches, and checks that each is 
remembered.
   
   On the testing point: not throwing a compile error of the unsplit code that 
the split avoids is intended. A "Code grows beyond 64 KB" failure that the 
split fixes is the case this PR exists for, and telling it apart from other 
failures would mean matching compiler messages. A test that needs the unsplit 
code's error can turn `spark.sql.codegen.wholeStage.splitExpressions` off.
   



##########
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:
   You're right, thanks. I made the kill switch exact rather than qualify it. 
Outside the whole-stage split policy, the slot rule now applies only when the 
context's `wholeStageSplit` is not `Off`. So with the conf off the code is 
master's again, including the slot passed as a parameter, and the description's 
two sentences are true. `CollectInputsSuite` has a new case, "with the whole 
stage split off, a slot is answered as before that split existed".
   



##########
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:
   Good catch, thank you. `CodeCompiler.updateMetrics` now defers a metric 
update while a trial is running. The class and method bytecode histograms in 
`computeByteCodeStats` go through it, as do the cache loader's source size, 
compile time and "Code generated in" log. `TrialCompile.report()` replays them 
before the huge-method report, and `chooseCode` calls it only for the code it 
keeps. So the counts stay at one per class that runs. The new test "only the 
code a stage keeps counts in the codegen metrics" checks that 
`METRIC_COMPILATION_TIME` grows by exactly one for a stage that splits. 
`DatasetOptimizationSuite` and `ArtifactManagerSuite` pass.
   



##########
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:
   You're right, thanks. It is now a `def df: DataFrame`, like the other tests.
   



##########
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:
   Agreed, thanks. The header scanner, `TypeKeyword`, `MethodHeader`, 
`wholeStageSplitSource` and `wholeStageSplitMarkersByMethod` are gone. The map 
is `CodegenContext.wholeStageSplitMarkersByFunction`, derived from 
`classFunctions` when the decision needs it, and computed only when some method 
is over the limit. It reflects only the functions actually declared, after 
`removeFunctionsSince` and replaced names, and it uses the literal-aware 
matcher from my reply on L2389.
   
   Separately, `removeFunctionsSince` subtracted the unstripped length from 
`classSize`, while `addNewFunctionToClass` added the stripped one. Both now go 
through one `sizeInClass`.
   



##########
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:
   Good point, thanks. `classLoaderKey()` is now typed `HashableWeakReference`, 
and one private `cacheKey(code)` is used by both `compile` and 
`invalidateCompiled`.
   



##########
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:
   Right, thanks. The parameter is removed from `splitExpressionsWithSources` 
and from the whole-stage helpers (`splitWholeStage`, `wholeStageBlocks`, 
`wholeStageSplitPossible`, `wholeStageSplitArguments`). A caller that needs it 
can add it back, along with a decision on how extra arguments relate to the 
audit.
   



##########
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:
   Thanks, agreed. The marker and call-site helpers now sit in a group of their 
own, with a section header, after the JVM-limit constants. With the scanner 
gone, the group is small.
   



-- 
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]

Reply via email to