dongjoon-hyun commented on code in PR #59069:
URL: https://github.com/apache/spark/pull/59069#discussion_r4118988121
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -1188,6 +1198,270 @@ class CodegenContext extends Logging {
}
}
+ /**
+ * `splitExpressionsWithCurrentInputs` for code whose sources 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, so `sources` gives those, one sequence for each entry of
`expressions`. Each function
+ * split out takes the locals its own entries read (see
`wholeStageSplitArguments`), and a
+ * function that groups calls takes what its calls take. The code stays in
one piece where the
+ * code of any function reads a local no function could take. Outside whole
stage codegen this is
+ * the method above.
+ *
+ * It is an overload rather than a new parameter of the method above, whose
signature code
+ * outside Spark is compiled against.
+ *
+ * @param expressions the codes to evaluate expressions.
+ * @param sources the expressions each entry of `expressions` was generated
from.
+ * @param funcName the split function name base.
+ * @param extraArguments the list of (type, name) of the arguments of the
split functions,
+ * except for 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.
+ */
+ def splitExpressionsWithCurrentInputs(
+ expressions: Seq[String],
+ sources: Seq[Seq[Expression]],
+ funcName: String,
+ extraArguments: Seq[(String, String)],
+ returnType: String,
+ makeSplitFunction: String => String,
+ foldFunctions: Seq[String] => String): String = {
+ require(sources.length == expressions.length,
+ s"${sources.length} source sequences for ${expressions.length} pieces of
code")
+ if (currentVars == null) {
+ splitExpressionsWithCurrentInputs(
+ expressions, funcName, extraArguments, returnType, makeSplitFunction,
foldFunctions)
+ } else if (!SQLConf.get.wholeStageSplitExpressions) {
+ expressions.mkString("\n")
+ } else {
+ splitWholeStage(expressions, sources, funcName, extraArguments,
returnType,
+ makeSplitFunction, foldFunctions).getOrElse(expressions.mkString("\n"))
+ }
+ }
+
+ /**
+ * The whole stage split of `splitExpressionsWithCurrentInputs`: None where
the code is one
+ * block, or where the code of a block reads a local no function could take.
+ */
+ private def splitWholeStage(
+ expressions: Seq[String],
+ sources: Seq[Seq[Expression]],
+ funcName: String,
+ extraArguments: Seq[(String, String)],
+ returnType: String,
+ makeSplitFunction: String => String,
+ foldFunctions: Seq[String] => String): Option[String] = {
+ val blocks = buildCodeBlocksOf(expressions)
+ if (blocks.length <= 1) {
+ None
+ } else {
+ val blockArguments = blocks.map { case (_, entries) =>
+ wholeStageSplitArguments(entries.flatMap(sources), extraArguments)
+ }
+ Option.when(blockArguments.forall(_.isDefined)) {
+ val func = freshName(funcName)
+ val calls = blocks.zip(blockArguments.flatten).zipWithIndex.map {
+ case (((body, _), arguments), i) =>
+ splitCall(s"${func}_$i", arguments ++ extraArguments, returnType,
+ makeSplitFunction(body))
+ }
+ foldFunctions(
+ groupSplitCalls(calls, func, returnType, makeSplitFunction,
foldFunctions).map(_.call))
+ }
+ }
+ }
+
+ /** A call to a function split out of whole stage generated code, and the
arguments it takes. */
+ private case class SplitCall(call: String, arguments: Seq[(String, String)])
+
+ /** Adds the function `name` taking `arguments` and running `body`, and
returns a call to it. */
+ private def splitCall(
+ name: String,
+ arguments: Seq[(String, String)],
+ returnType: String,
+ body: String): SplitCall = {
+ if (Utils.isTesting) {
+ // Passing global variables to the split method is dangerous, as any
mutating to it is
+ // ignored and may lead to unexpected behavior.
+ arguments.foreach { case (_, argument) =>
+ assert(!mutableStateNames.contains(argument),
+ s"split function argument $argument cannot be a global variable.")
+ }
+ }
+ val argString = arguments.map { case (t, argument) => s"$t $argument"
}.mkString(", ")
+ val code =
+ s"""
+ |private $returnType $name($argString) {
+ | $body
+ |}
+ """.stripMargin
+ SplitCall(s"${addNewFunction(name,
code)}(${arguments.map(_._2).mkString(", ")})", arguments)
+ }
+
+ /**
+ * `groupFunctionCalls` for calls that take different arguments: a group
takes the union of what
+ * its calls take, and is closed when one more call would take it past
+ * `SPLIT_CALLS_PER_GROUP` calls or past the parameters a method can have.
Where no group can
+ * hold two calls, the calls are left as they are.
+ */
+ private def groupSplitCalls(
+ calls: Seq[SplitCall],
+ funcName: String,
+ returnType: String,
+ makeSplitFunction: String => String,
+ foldFunctions: Seq[String] => String,
+ level: Int = 0): Seq[SplitCall] = {
+ if (calls.isEmpty ||
Review Comment:
`groupSplitCalls` leaves up to `splitCallsPerMethod` calls per CASE WHEN
ungrouped (~166 for one nullable input), but under whole-stage codegen the call
sites of all CASE WHENs of an operator land in the same method. For instance,
the outputs of `ProjectExec` are concatenated by `evaluateVariables(colVars)`
in `prepareRowVar`.
With 4 CASE WHENs of 300 branches (~2.5KB of call sites each, see `2633
bytes split` in the benchmark results) or ~14 of 64 branches in one projection,
that method is ~10KB:
- above HotSpot's 8000-byte limit, so it's never JIT-compiled, and
- below `spark.sql.codegen.hugeMethodLimit` (65535), so there's no fallback
either.
Before this PR, the same stage exceeded 64KB, failed to compile and fell
back to non-whole-stage code, where each CASE WHEN's calls sit in their own
JIT-compiled method. So this shape may get slower (I haven't measured it).
Could you check, e.g., four 300-branch CASE WHENs over different columns,
default vs `spark.sql.codegen.wholeStage=false`? Always grouping above
`SPLIT_CALLS_PER_GROUP` under whole-stage codegen, or tracking the budget per
method, would avoid it.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -1188,6 +1198,270 @@ class CodegenContext extends Logging {
}
}
+ /**
+ * `splitExpressionsWithCurrentInputs` for code whose sources 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, so `sources` gives those, one sequence for each entry of
`expressions`. Each function
+ * split out takes the locals its own entries read (see
`wholeStageSplitArguments`), and a
+ * function that groups calls takes what its calls take. The code stays in
one piece where the
+ * code of any function reads a local no function could take. Outside whole
stage codegen this is
+ * the method above.
+ *
+ * It is an overload rather than a new parameter of the method above, whose
signature code
+ * outside Spark is compiled against.
+ *
+ * @param expressions the codes to evaluate expressions.
+ * @param sources the expressions each entry of `expressions` was generated
from.
+ * @param funcName the split function name base.
+ * @param extraArguments the list of (type, name) of the arguments of the
split functions,
+ * except for 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.
+ */
+ def splitExpressionsWithCurrentInputs(
+ expressions: Seq[String],
+ sources: Seq[Seq[Expression]],
+ funcName: String,
+ extraArguments: Seq[(String, String)],
+ returnType: String,
+ makeSplitFunction: String => String,
+ foldFunctions: Seq[String] => String): String = {
+ require(sources.length == expressions.length,
+ s"${sources.length} source sequences for ${expressions.length} pieces of
code")
+ if (currentVars == null) {
+ splitExpressionsWithCurrentInputs(
+ expressions, funcName, extraArguments, returnType, makeSplitFunction,
foldFunctions)
+ } else if (!SQLConf.get.wholeStageSplitExpressions) {
+ expressions.mkString("\n")
+ } else {
+ splitWholeStage(expressions, sources, funcName, extraArguments,
returnType,
+ makeSplitFunction, foldFunctions).getOrElse(expressions.mkString("\n"))
+ }
+ }
+
+ /**
+ * The whole stage split of `splitExpressionsWithCurrentInputs`: None where
the code is one
+ * block, or where the code of a block reads a local no function could take.
+ */
+ private def splitWholeStage(
+ expressions: Seq[String],
+ sources: Seq[Seq[Expression]],
+ funcName: String,
+ extraArguments: Seq[(String, String)],
+ returnType: String,
+ makeSplitFunction: String => String,
+ foldFunctions: Seq[String] => String): Option[String] = {
+ val blocks = buildCodeBlocksOf(expressions)
+ if (blocks.length <= 1) {
+ None
+ } else {
+ val blockArguments = blocks.map { case (_, entries) =>
+ wholeStageSplitArguments(entries.flatMap(sources), extraArguments)
+ }
+ Option.when(blockArguments.forall(_.isDefined)) {
Review Comment:
Since one refused block keeps the whole CASE WHEN inline, common shapes
still aren't split. For example, `SELECT CASE code WHEN 1 THEN 'a' ... WHEN N
THEN '...' ELSE name END` over a Parquet table:
- `code` is read N times, so `ProjectExec` evaluates it up front.
- `name` is read once, so `ProjectExec` defers it (`input.code !=
EmptyBlock`), and the ELSE block is refused.
- The entire CASE WHEN stays in the stage method, so a large one is still
past 8000 bytes or 64KB, as before this PR.
The fuzz run in the PR description also reports 430 refusals, all for this
reason. How about a per-block fallback that keeps only the refused blocks
inline between the calls? `CaseWhen`'s fold is per call and shares the same `do
{ } while (false)` and `resultState`, so they can be interleaved (grouping only
consecutive calls). The pivot shape (`WHEN k = i THEN c_i`) would still need
something else.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/CaseWhenCodegenBenchmark.scala:
##########
@@ -0,0 +1,95 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.benchmark
+
+import org.apache.spark.benchmark.Benchmark
+import org.apache.spark.sql.DataFrame
+import org.apache.spark.sql.catalyst.expressions.codegen.CodeGenerator
+import org.apache.spark.sql.execution.WholeStageCodegenExec
+import org.apache.spark.sql.internal.SQLConf
+
+/**
+ * A `CASE WHEN` of 2 to 300 branches in a projection, with whole-stage
codegen splitting its
+ * code into methods (`spark.sql.codegen.wholeStage.splitExpressions`, the
default), without the
+ * split, and without whole-stage codegen.
+ *
+ * The split takes effect once the branches' code passes
`spark.sql.codegen.methodSplitThreshold`
+ * characters, a few branches in, so the small rungs price a call per method
against code kept in
+ * the stage's method, and the large ones the stage's method staying under the
8000 bytes HotSpot
+ * compiles against one past it. Each rung prints the largest method of the
stage with the split
+ * and without it above its timings.
+ *
+ * To run this benchmark:
+ * {{{
+ * 1. without sbt:
+ * bin/spark-submit --class <this class> --jars <spark core test jar>
<spark sql test jar>
+ * 2. build/sbt "sql/Test/runMain <this class>"
+ * 3. generate result:
+ * SPARK_GENERATE_BENCHMARK_FILES=1 build/sbt "sql/Test/runMain <this
class>"
+ * Results will be written to
"benchmarks/CaseWhenCodegenBenchmark-results.txt".
+ * }}}
+ */
+object CaseWhenCodegenBenchmark extends SqlBasedBenchmark {
+
+ private val rows = 2L * 1000 * 1000
+
+ private def query(branches: Int): DataFrame = {
+ val whens = (1 to branches).map(k => s"WHEN v = $k THEN v *
$k").mkString(" ")
+ spark.sql(
+ s"SELECT CASE $whens ELSE 0 END AS c FROM (SELECT id % $branches AS v
FROM range($rows))")
+ }
+
+ /** The largest method of the whole-stage codegen stage of `df`'s plan,
compiled. */
+ private def stageMethodSize(df: DataFrame): Int =
+ df.queryExecution.executedPlan.collect { case w: WholeStageCodegenExec =>
w }
+ .map(s =>
CodeGenerator.compile(s.doCodeGen()._2)._2.maxMethodCodeSize).max
+
+ private def caseWhen(branches: Int): Unit = {
+ val benchmark = new Benchmark(s"CASE WHEN of $branches branches", rows,
output = output)
+ val split = stageMethodSize(query(branches))
+ val unsplit = withSQLConf(SQLConf.WHOLESTAGE_SPLIT_EXPRESSIONS.key ->
"false") {
+ stageMethodSize(query(branches))
+ }
+ // scalastyle:off println
+ benchmark.out.println(
+ s"largest method of the stage: $split bytes split, $unsplit bytes not
split")
+ // scalastyle:on println
+ benchmark.addCase("whole-stage codegen, split", numIters = 5) { _ =>
Review Comment:
I'm not sure this benchmark can resolve the differences the PR description
relies on:
- There's no real A/A control. At 2 branches the ELSE is already split (a
branch with ANSI arithmetic is ~740 chars). The `280 bytes` printed for both
variants seems to be another method of the stage, not the CASE WHEN.
- The case order is fixed (the split case always runs first), and every
iteration re-parses, re-plans and re-generates code inside the timed lambda.
- The differences called noise (2 branches: +27.5/-10.8/+20.0%, 32 branches:
+4.5/-14.8/-16.5% on JDK 17/21/25) are as large as the 64-branch regression
(+12.6/+21.2/+18.8%).
- There's only one input shape (a single long column).
Could we add a same-conf A/A pair, alternate the order, build the DataFrame
once, and add shapes such as several CASE WHENs in one projection, string
branches and an aggregate? Also, the results show up to ~28x (128 branches, JDK
17), a bit beyond the 14-26x in the description.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -1562,6 +1854,9 @@ class CodegenContext extends Logging {
val needSplit = nonSplitCode.map(_.eval.code.length).sum >
SQLConf.get.methodSplitThreshold
val (subExprsMap, exprCodes) = if (needSplit) {
if
(inputVarsForAllFuncs.map(calculateParamLengthFromExprValues).forall(isValidParamLength))
{
+ nonSplitFunctions.foreach { case (className, funcName) =>
Review Comment:
`spark.sql.codegen.wholeStage.splitExpressions` only disables the CASE WHEN
split (L1236). The recording and this removal are unconditional. They apply to
every whole-stage subexpression-elimination split pass and to functions of any
expression (`computeCommonExpr`, `compare*` from `genComp`, `ToStringBase`
helpers, `Concat`, ...), not only CASE WHEN.
I couldn't find an in-tree case where surviving code still references such a
function. But if one appears, e.g. an expression that caches a helper name or
calls it from mutable-state init code, setting the conf to false won't restore
the previous behavior. Could we gate the removal with the same conf? With the
conf off there are no CASE WHEN split functions to remove anyway.
Alternatively, the conf doc could state its scope.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -1188,6 +1198,270 @@ class CodegenContext extends Logging {
}
}
+ /**
+ * `splitExpressionsWithCurrentInputs` for code whose sources 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, so `sources` gives those, one sequence for each entry of
`expressions`. Each function
+ * split out takes the locals its own entries read (see
`wholeStageSplitArguments`), and a
+ * function that groups calls takes what its calls take. The code stays in
one piece where the
+ * code of any function reads a local no function could take. Outside whole
stage codegen this is
+ * the method above.
+ *
+ * It is an overload rather than a new parameter of the method above, whose
signature code
+ * outside Spark is compiled against.
+ *
+ * @param expressions the codes to evaluate expressions.
+ * @param sources the expressions each entry of `expressions` was generated
from.
+ * @param funcName the split function name base.
+ * @param extraArguments the list of (type, name) of the arguments of the
split functions,
+ * except for 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.
+ */
+ def splitExpressionsWithCurrentInputs(
+ expressions: Seq[String],
+ sources: Seq[Seq[Expression]],
+ funcName: String,
+ extraArguments: Seq[(String, String)],
+ returnType: String,
+ makeSplitFunction: String => String,
+ foldFunctions: Seq[String] => String): String = {
+ require(sources.length == expressions.length,
+ s"${sources.length} source sequences for ${expressions.length} pieces of
code")
+ if (currentVars == null) {
+ splitExpressionsWithCurrentInputs(
+ expressions, funcName, extraArguments, returnType, makeSplitFunction,
foldFunctions)
+ } else if (!SQLConf.get.wholeStageSplitExpressions) {
+ expressions.mkString("\n")
+ } else {
+ splitWholeStage(expressions, sources, funcName, extraArguments,
returnType,
+ makeSplitFunction, foldFunctions).getOrElse(expressions.mkString("\n"))
+ }
+ }
+
+ /**
+ * The whole stage split of `splitExpressionsWithCurrentInputs`: None where
the code is one
+ * block, or where the code of a block reads a local no function could take.
+ */
+ private def splitWholeStage(
+ expressions: Seq[String],
+ sources: Seq[Seq[Expression]],
+ funcName: String,
+ extraArguments: Seq[(String, String)],
+ returnType: String,
+ makeSplitFunction: String => String,
+ foldFunctions: Seq[String] => String): Option[String] = {
+ val blocks = buildCodeBlocksOf(expressions)
+ if (blocks.length <= 1) {
+ None
+ } else {
+ val blockArguments = blocks.map { case (_, entries) =>
+ wholeStageSplitArguments(entries.flatMap(sources), extraArguments)
+ }
+ Option.when(blockArguments.forall(_.isDefined)) {
+ val func = freshName(funcName)
+ val calls = blocks.zip(blockArguments.flatten).zipWithIndex.map {
+ case (((body, _), arguments), i) =>
+ splitCall(s"${func}_$i", arguments ++ extraArguments, returnType,
+ makeSplitFunction(body))
+ }
+ foldFunctions(
+ groupSplitCalls(calls, func, returnType, makeSplitFunction,
foldFunctions).map(_.call))
+ }
+ }
+ }
+
+ /** A call to a function split out of whole stage generated code, and the
arguments it takes. */
+ private case class SplitCall(call: String, arguments: Seq[(String, String)])
+
+ /** Adds the function `name` taking `arguments` and running `body`, and
returns a call to it. */
+ private def splitCall(
+ name: String,
+ arguments: Seq[(String, String)],
+ returnType: String,
+ body: String): SplitCall = {
+ if (Utils.isTesting) {
+ // Passing global variables to the split method is dangerous, as any
mutating to it is
+ // ignored and may lead to unexpected behavior.
+ arguments.foreach { case (_, argument) =>
+ assert(!mutableStateNames.contains(argument),
+ s"split function argument $argument cannot be a global variable.")
+ }
+ }
+ val argString = arguments.map { case (t, argument) => s"$t $argument"
}.mkString(", ")
+ val code =
+ s"""
+ |private $returnType $name($argString) {
+ | $body
+ |}
+ """.stripMargin
+ SplitCall(s"${addNewFunction(name,
code)}(${arguments.map(_._2).mkString(", ")})", arguments)
+ }
+
+ /**
+ * `groupFunctionCalls` for calls that take different arguments: a group
takes the union of what
+ * its calls take, and is closed when one more call would take it past
+ * `SPLIT_CALLS_PER_GROUP` calls or past the parameters a method can have.
Where no group can
+ * hold two calls, the calls are left as they are.
+ */
+ private def groupSplitCalls(
+ calls: Seq[SplitCall],
+ funcName: String,
+ returnType: String,
+ makeSplitFunction: String => String,
+ foldFunctions: Seq[String] => String,
+ level: Int = 0): Seq[SplitCall] = {
+ if (calls.isEmpty ||
+ calls.length <=
splitCallsPerMethod(calls.maxBy(_.arguments.length).arguments)) {
+ calls
+ } else {
+ val groups = mutable.ArrayBuffer.empty[mutable.ArrayBuffer[SplitCall]]
+ var taken = Seq.empty[(String, String)]
+ calls.foreach { c =>
+ val union = (taken ++ c.arguments).distinctBy(_._2)
+ if (groups.isEmpty || groups.last.length == SPLIT_CALLS_PER_GROUP ||
+ !isValidParamLength(paramSlots(union))) {
+ groups += mutable.ArrayBuffer(c)
+ taken = c.arguments
+ } else {
+ groups.last += c
+ taken = union
+ }
+ }
+ if (groups.length == calls.length) {
+ calls
+ } else {
+ val grouped = groups.zipWithIndex.map { case (group, i) =>
+ splitCall(s"${funcName}_group${level}_$i",
+ group.flatMap(_.arguments).distinctBy(_._2).toSeq, returnType,
+ makeSplitFunction(foldFunctions(group.map(_.call).toSeq)))
+ }.toSeq
+ groupSplitCalls(grouped, funcName, returnType, makeSplitFunction,
foldFunctions, level + 1)
+ }
+ }
+ }
+
+ /** The parameter slots of a method taking `arguments`, `this` included. */
+ private def paramSlots(arguments: Seq[(String, String)]): Int =
+ 1 + arguments.map { case (javaType, _) =>
+ if (javaType == JAVA_LONG || javaType == JAVA_DOUBLE) 2 else 1
+ }.sum
+
+ /**
+ * The arguments of a function split out of whole stage generated code so
that the code generated
+ * for `sources` compiles inside it: the local variables that code reads
from the operator's
+ * method. None where the code cannot be split, because a local it reads
cannot be told from
+ * `sources` or cannot be passed.
+ *
+ * The code reads its inputs from `currentVars`, from the states of
subexpression elimination
+ * and, where the operator sets one, from `INPUT_ROW`; a field or a literal
among them is read as
+ * it stands and needs no argument. A state is taken as its own variables
without looking below
+ * it, the way `CodeGenerator.getLocalInputVariableValues` takes it, so that
the arguments are
+ * among those of the functions an aggregate or `ExpandExec` splits out of
its own code, which is
+ * where these calls then sit. A `CommonExpressionRef` computes its
definition where it is
+ * reached, so what the definition reads is read here too.
+ *
+ * It refuses, as `CommonExprSlots.methodArgs` does and for the reasons it
gives:
+ * - an input the operator has not evaluated yet, whose code was generated
into this code and
+ * names locals of the operator producing the row;
+ * - a value no parameter can carry: a `SimpleExprValue`, or a slot of a
compacted mutable
+ * state array such as `mutableStateArray_0[3]`, which is not a name;
+ * - more parameter slots than a method can take, `extraArguments` included.
+ * And two it has of its own:
+ * - a state for a node that overrides `genCode` and so may generate its
child instead of reading
+ * the state (`Alias`, `Collate` and a `Cast` that changes no structure
do), so that what the
+ * code reads below it is not in the state;
+ * - a `LambdaVariable`. Its value is a field that the loop of its object
expression assigns,
+ * which a function could read as it is, but no test reaches such a
split, so the code keeps
+ * its place.
+ * `INPUT_ROW` is passed only when the code reads it: through an input the
operator has not put
+ * in `currentVars`, or through a `CommonExpressionRef`, whose definition
`CommonExprSlots` may
+ * make a method that takes `INPUT_ROW` wherever it is set. An operator may
leave `INPUT_ROW`
+ * naming a local that is not in scope where other code is, as a join does
after generating the
+ * variables of its build side, which is why it is not passed otherwise.
+ */
+ private def wholeStageSplitArguments(
+ sources: Seq[Expression],
+ extraArguments: Seq[(String, String)]): Option[Seq[(String, String)]] = {
+ val args = mutable.LinkedHashMap.empty[String, VariableValue]
+ def isName(name: String): Boolean = name.nonEmpty &&
+ Character.isJavaIdentifierStart(name.head) &&
name.forall(Character.isJavaIdentifierPart)
+ def canPass(v: ExprValue): Boolean = v match {
+ case local: VariableValue =>
+ val name = local.variableName
+ // A field needs no argument, and passing one is refused by
`splitExpressions`.
+ if (isName(name) && !mutableStateNames.contains(name)) {
+ args.getOrElseUpdate(name, local)
+ }
+ isName(name)
+ case _: GlobalValue | _: LiteralValue => true
+ case _ => false
+ }
+ // `Expression.genCode` is what reads a state; a node that overrides it
may not.
+ def readsSubExprState(e: Expression): Boolean = e match {
+ case c: Cast => !c.generatesChildCode
+ case _ => e.getClass.getMethod("genCode",
classOf[CodegenContext]).getDeclaringClass ==
+ classOf[Expression]
+ }
+ var needsRow = false
+ var possible = sources.nonEmpty
+ val visited = mutable.HashSet.empty[Long]
+ // The definitions of a `With` in the trees walked here, which are in
`currentCommonExprs`
+ // only while that `With` is generated.
+ val nestedDefs = mutable.HashMap.empty[Long, Expression]
+ val toVisit = mutable.Stack[Expression](sources: _*)
+ while (possible && toVisit.nonEmpty) {
+ val next = toVisit.pop()
+ subExprEliminationExprs.get(ExpressionEquals(next)) match {
+ case Some(state) =>
+ possible = readsSubExprState(next) &&
+ canPass(state.eval.value) && canPass(state.eval.isNull)
+ case None =>
+ next match {
+ case ref: BoundReference if currentVars(ref.ordinal) != null =>
+ val input = currentVars(ref.ordinal)
+ possible = input.code == EmptyBlock && canPass(input.value) &&
canPass(input.isNull)
+ case _: BoundReference =>
+ needsRow = true
+ case _: LambdaVariable =>
+ possible = false
+ case w: With =>
+ w.defs.foreach(d => nestedDefs.put(d.id.id, d.child))
+ toVisit.push(w.child)
+ case ref: CommonExpressionRef =>
+ // The call to a definition made a method passes `INPUT_ROW`
wherever it is set, and
+ // whether the definition is made one depends on the length of
its code
+ // (`CommonExprSlots.build`), which is generated only after
these arguments are
+ // computed. So the row is taken wherever it is set: for a
definition that ends up
+ // inlined, an argument the method does not read. That is safe
where references
+ // reach code generation - an aggregate, where `INPUT_ROW` is in
scope -
+ // `RewriteWithExpression` having inlined a `With` in a join
condition.
+ needsRow |= INPUT_ROW != null
Review Comment:
The comment above doesn't seem to match the code:
- `RewriteWithExpression` keeps a `With` in a non-first conditional branch
for every operator, joins included (`inlineDefsThatGainNothing`), e.g. `... ON
l.k = r.k AND CASE WHEN l.v = 0 THEN 0 WHEN l.v = 1 THEN nullif(l.v * 3 + r.w,
1) ... END > 50`.
- `CaseWhen` generates all branches before calling the split, so
`CommonExprSlots.fill` has already decided between a method and inlining. The
walk just can't observe that decision.
So passing `INPUT_ROW` for every reference is safe only because `INPUT_ROW`
happens to be in scope wherever such a CASE WHEN is generated today. That
doesn't hold at the streamed-only condition of `HashJoin`
(outer/anti/existence) and `SortMergeJoinExec`: there `INPUT_ROW` names
`matched`/`bufferedRow`, which is declared after that code. Today it's masked
only by `ExprUtils.canEvaluateUnconditionally` (no `CaseWhen`/`With`) and
`spark.sql.join.splitStreamedSideJoinCondition=false`. Could we fix the
comment, or record in the slot whether the definition became a method and pass
the row only then?
Also, the LEFT JOIN query in the new join test doesn't enable
`spark.sql.join.splitStreamedSideJoinCondition`, so it covers the
rest-condition path rather than the streamed-only condition mentioned in the PR
description.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -1188,6 +1198,270 @@ class CodegenContext extends Logging {
}
}
+ /**
+ * `splitExpressionsWithCurrentInputs` for code whose sources 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, so `sources` gives those, one sequence for each entry of
`expressions`. Each function
+ * split out takes the locals its own entries read (see
`wholeStageSplitArguments`), and a
+ * function that groups calls takes what its calls take. The code stays in
one piece where the
+ * code of any function reads a local no function could take. Outside whole
stage codegen this is
+ * the method above.
+ *
+ * It is an overload rather than a new parameter of the method above, whose
signature code
+ * outside Spark is compiled against.
+ *
+ * @param expressions the codes to evaluate expressions.
+ * @param sources the expressions each entry of `expressions` was generated
from.
+ * @param funcName the split function name base.
+ * @param extraArguments the list of (type, name) of the arguments of the
split functions,
+ * except for 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.
+ */
+ def splitExpressionsWithCurrentInputs(
+ expressions: Seq[String],
+ sources: Seq[Seq[Expression]],
+ funcName: String,
+ extraArguments: Seq[(String, String)],
+ returnType: String,
+ makeSplitFunction: String => String,
+ foldFunctions: Seq[String] => String): String = {
+ require(sources.length == expressions.length,
+ s"${sources.length} source sequences for ${expressions.length} pieces of
code")
+ if (currentVars == null) {
+ splitExpressionsWithCurrentInputs(
+ expressions, funcName, extraArguments, returnType, makeSplitFunction,
foldFunctions)
+ } else if (!SQLConf.get.wholeStageSplitExpressions) {
+ expressions.mkString("\n")
+ } else {
+ splitWholeStage(expressions, sources, funcName, extraArguments,
returnType,
+ makeSplitFunction, foldFunctions).getOrElse(expressions.mkString("\n"))
+ }
+ }
+
+ /**
+ * The whole stage split of `splitExpressionsWithCurrentInputs`: None where
the code is one
+ * block, or where the code of a block reads a local no function could take.
+ */
+ private def splitWholeStage(
+ expressions: Seq[String],
+ sources: Seq[Seq[Expression]],
+ funcName: String,
+ extraArguments: Seq[(String, String)],
+ returnType: String,
+ makeSplitFunction: String => String,
+ foldFunctions: Seq[String] => String): Option[String] = {
+ val blocks = buildCodeBlocksOf(expressions)
+ if (blocks.length <= 1) {
+ None
+ } else {
+ val blockArguments = blocks.map { case (_, entries) =>
+ wholeStageSplitArguments(entries.flatMap(sources), extraArguments)
+ }
+ Option.when(blockArguments.forall(_.isDefined)) {
+ val func = freshName(funcName)
+ val calls = blocks.zip(blockArguments.flatten).zipWithIndex.map {
+ case (((body, _), arguments), i) =>
+ splitCall(s"${func}_$i", arguments ++ extraArguments, returnType,
+ makeSplitFunction(body))
+ }
+ foldFunctions(
+ groupSplitCalls(calls, func, returnType, makeSplitFunction,
foldFunctions).map(_.call))
+ }
+ }
+ }
+
+ /** A call to a function split out of whole stage generated code, and the
arguments it takes. */
+ private case class SplitCall(call: String, arguments: Seq[(String, String)])
+
+ /** Adds the function `name` taking `arguments` and running `body`, and
returns a call to it. */
+ private def splitCall(
+ name: String,
+ arguments: Seq[(String, String)],
+ returnType: String,
+ body: String): SplitCall = {
+ if (Utils.isTesting) {
+ // Passing global variables to the split method is dangerous, as any
mutating to it is
+ // ignored and may lead to unexpected behavior.
+ arguments.foreach { case (_, argument) =>
+ assert(!mutableStateNames.contains(argument),
+ s"split function argument $argument cannot be a global variable.")
+ }
+ }
+ val argString = arguments.map { case (t, argument) => s"$t $argument"
}.mkString(", ")
+ val code =
+ s"""
+ |private $returnType $name($argString) {
+ | $body
+ |}
+ """.stripMargin
+ SplitCall(s"${addNewFunction(name,
code)}(${arguments.map(_._2).mkString(", ")})", arguments)
+ }
+
+ /**
+ * `groupFunctionCalls` for calls that take different arguments: a group
takes the union of what
+ * its calls take, and is closed when one more call would take it past
+ * `SPLIT_CALLS_PER_GROUP` calls or past the parameters a method can have.
Where no group can
+ * hold two calls, the calls are left as they are.
+ */
+ private def groupSplitCalls(
+ calls: Seq[SplitCall],
+ funcName: String,
+ returnType: String,
+ makeSplitFunction: String => String,
+ foldFunctions: Seq[String] => String,
+ level: Int = 0): Seq[SplitCall] = {
+ if (calls.isEmpty ||
+ calls.length <=
splitCallsPerMethod(calls.maxBy(_.arguments.length).arguments)) {
+ calls
+ } else {
+ val groups = mutable.ArrayBuffer.empty[mutable.ArrayBuffer[SplitCall]]
+ var taken = Seq.empty[(String, String)]
+ calls.foreach { c =>
+ val union = (taken ++ c.arguments).distinctBy(_._2)
+ if (groups.isEmpty || groups.last.length == SPLIT_CALLS_PER_GROUP ||
+ !isValidParamLength(paramSlots(union))) {
+ groups += mutable.ArrayBuffer(c)
+ taken = c.arguments
+ } else {
+ groups.last += c
+ taken = union
+ }
+ }
+ if (groups.length == calls.length) {
+ calls
+ } else {
+ val grouped = groups.zipWithIndex.map { case (group, i) =>
+ splitCall(s"${funcName}_group${level}_$i",
+ group.flatMap(_.arguments).distinctBy(_._2).toSeq, returnType,
+ makeSplitFunction(foldFunctions(group.map(_.call).toSeq)))
+ }.toSeq
+ groupSplitCalls(grouped, funcName, returnType, makeSplitFunction,
foldFunctions, level + 1)
+ }
+ }
+ }
+
+ /** The parameter slots of a method taking `arguments`, `this` included. */
+ private def paramSlots(arguments: Seq[(String, String)]): Int =
+ 1 + arguments.map { case (javaType, _) =>
+ if (javaType == JAVA_LONG || javaType == JAVA_DOUBLE) 2 else 1
+ }.sum
+
+ /**
+ * The arguments of a function split out of whole stage generated code so
that the code generated
+ * for `sources` compiles inside it: the local variables that code reads
from the operator's
+ * method. None where the code cannot be split, because a local it reads
cannot be told from
+ * `sources` or cannot be passed.
+ *
+ * The code reads its inputs from `currentVars`, from the states of
subexpression elimination
+ * and, where the operator sets one, from `INPUT_ROW`; a field or a literal
among them is read as
+ * it stands and needs no argument. A state is taken as its own variables
without looking below
+ * it, the way `CodeGenerator.getLocalInputVariableValues` takes it, so that
the arguments are
+ * among those of the functions an aggregate or `ExpandExec` splits out of
its own code, which is
+ * where these calls then sit. A `CommonExpressionRef` computes its
definition where it is
+ * reached, so what the definition reads is read here too.
+ *
+ * It refuses, as `CommonExprSlots.methodArgs` does and for the reasons it
gives:
+ * - an input the operator has not evaluated yet, whose code was generated
into this code and
+ * names locals of the operator producing the row;
+ * - a value no parameter can carry: a `SimpleExprValue`, or a slot of a
compacted mutable
+ * state array such as `mutableStateArray_0[3]`, which is not a name;
+ * - more parameter slots than a method can take, `extraArguments` included.
+ * And two it has of its own:
+ * - a state for a node that overrides `genCode` and so may generate its
child instead of reading
+ * the state (`Alias`, `Collate` and a `Cast` that changes no structure
do), so that what the
+ * code reads below it is not in the state;
+ * - a `LambdaVariable`. Its value is a field that the loop of its object
expression assigns,
+ * which a function could read as it is, but no test reaches such a
split, so the code keeps
+ * its place.
+ * `INPUT_ROW` is passed only when the code reads it: through an input the
operator has not put
+ * in `currentVars`, or through a `CommonExpressionRef`, whose definition
`CommonExprSlots` may
+ * make a method that takes `INPUT_ROW` wherever it is set. An operator may
leave `INPUT_ROW`
+ * naming a local that is not in scope where other code is, as a join does
after generating the
+ * variables of its build side, which is why it is not passed otherwise.
+ */
+ private def wholeStageSplitArguments(
Review Comment:
`wholeStageSplitArguments` is mostly a copy of `CommonExprSlots.methodArgs`
above: the `args` map, `isName`/`canPass`, `visited`/`nestedDefs`/`toVisit`,
the `BoundReference`/`With`/`CommonExpressionRef` arms and the final length
check. It's now the third input collector next to
`getLocalInputVariableValues`, and the three already have different policies
for:
- fields (`mutableStateNames`),
- `INPUT_ROW` (always vs. only when read),
- `LambdaVariable`, and
- subexpression states.
A nested split is correct only if this result is a subset of the parameters
that `getLocalInputVariableValues` computes for the enclosing
`subExpr_N`/`doAggregate_*`/`switchCaseCode`, and only the comment at
L1360-1363 guarantees that. Could we share one walker parameterized by these
policies, so that a new variable-binding node needs one change instead of three?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -1188,6 +1198,270 @@ class CodegenContext extends Logging {
}
}
+ /**
+ * `splitExpressionsWithCurrentInputs` for code whose sources 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, so `sources` gives those, one sequence for each entry of
`expressions`. Each function
+ * split out takes the locals its own entries read (see
`wholeStageSplitArguments`), and a
+ * function that groups calls takes what its calls take. The code stays in
one piece where the
+ * code of any function reads a local no function could take. Outside whole
stage codegen this is
+ * the method above.
+ *
+ * It is an overload rather than a new parameter of the method above, whose
signature code
+ * outside Spark is compiled against.
+ *
+ * @param expressions the codes to evaluate expressions.
+ * @param sources the expressions each entry of `expressions` was generated
from.
+ * @param funcName the split function name base.
+ * @param extraArguments the list of (type, name) of the arguments of the
split functions,
+ * except for 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.
+ */
+ def splitExpressionsWithCurrentInputs(
+ expressions: Seq[String],
+ sources: Seq[Seq[Expression]],
+ funcName: String,
+ extraArguments: Seq[(String, String)],
+ returnType: String,
+ makeSplitFunction: String => String,
+ foldFunctions: Seq[String] => String): String = {
+ require(sources.length == expressions.length,
+ s"${sources.length} source sequences for ${expressions.length} pieces of
code")
+ if (currentVars == null) {
+ splitExpressionsWithCurrentInputs(
+ expressions, funcName, extraArguments, returnType, makeSplitFunction,
foldFunctions)
+ } else if (!SQLConf.get.wholeStageSplitExpressions) {
+ expressions.mkString("\n")
+ } else {
+ splitWholeStage(expressions, sources, funcName, extraArguments,
returnType,
+ makeSplitFunction, foldFunctions).getOrElse(expressions.mkString("\n"))
+ }
+ }
+
+ /**
+ * The whole stage split of `splitExpressionsWithCurrentInputs`: None where
the code is one
+ * block, or where the code of a block reads a local no function could take.
+ */
+ private def splitWholeStage(
+ expressions: Seq[String],
+ sources: Seq[Seq[Expression]],
+ funcName: String,
+ extraArguments: Seq[(String, String)],
+ returnType: String,
+ makeSplitFunction: String => String,
+ foldFunctions: Seq[String] => String): Option[String] = {
+ val blocks = buildCodeBlocksOf(expressions)
+ if (blocks.length <= 1) {
+ None
+ } else {
+ val blockArguments = blocks.map { case (_, entries) =>
+ wholeStageSplitArguments(entries.flatMap(sources), extraArguments)
+ }
+ Option.when(blockArguments.forall(_.isDefined)) {
+ val func = freshName(funcName)
+ val calls = blocks.zip(blockArguments.flatten).zipWithIndex.map {
+ case (((body, _), arguments), i) =>
+ splitCall(s"${func}_$i", arguments ++ extraArguments, returnType,
+ makeSplitFunction(body))
+ }
+ foldFunctions(
+ groupSplitCalls(calls, func, returnType, makeSplitFunction,
foldFunctions).map(_.call))
+ }
+ }
+ }
+
+ /** A call to a function split out of whole stage generated code, and the
arguments it takes. */
+ private case class SplitCall(call: String, arguments: Seq[(String, String)])
+
+ /** Adds the function `name` taking `arguments` and running `body`, and
returns a call to it. */
+ private def splitCall(
+ name: String,
+ arguments: Seq[(String, String)],
+ returnType: String,
+ body: String): SplitCall = {
+ if (Utils.isTesting) {
+ // Passing global variables to the split method is dangerous, as any
mutating to it is
+ // ignored and may lead to unexpected behavior.
+ arguments.foreach { case (_, argument) =>
+ assert(!mutableStateNames.contains(argument),
+ s"split function argument $argument cannot be a global variable.")
+ }
+ }
+ val argString = arguments.map { case (t, argument) => s"$t $argument"
}.mkString(", ")
+ val code =
+ s"""
+ |private $returnType $name($argString) {
+ | $body
+ |}
+ """.stripMargin
+ SplitCall(s"${addNewFunction(name,
code)}(${arguments.map(_._2).mkString(", ")})", arguments)
+ }
+
+ /**
+ * `groupFunctionCalls` for calls that take different arguments: a group
takes the union of what
+ * its calls take, and is closed when one more call would take it past
+ * `SPLIT_CALLS_PER_GROUP` calls or past the parameters a method can have.
Where no group can
+ * hold two calls, the calls are left as they are.
+ */
+ private def groupSplitCalls(
+ calls: Seq[SplitCall],
+ funcName: String,
+ returnType: String,
+ makeSplitFunction: String => String,
+ foldFunctions: Seq[String] => String,
+ level: Int = 0): Seq[SplitCall] = {
+ if (calls.isEmpty ||
+ calls.length <=
splitCallsPerMethod(calls.maxBy(_.arguments.length).arguments)) {
+ calls
+ } else {
+ val groups = mutable.ArrayBuffer.empty[mutable.ArrayBuffer[SplitCall]]
+ var taken = Seq.empty[(String, String)]
+ calls.foreach { c =>
+ val union = (taken ++ c.arguments).distinctBy(_._2)
+ if (groups.isEmpty || groups.last.length == SPLIT_CALLS_PER_GROUP ||
+ !isValidParamLength(paramSlots(union))) {
+ groups += mutable.ArrayBuffer(c)
+ taken = c.arguments
+ } else {
+ groups.last += c
+ taken = union
+ }
+ }
+ if (groups.length == calls.length) {
+ calls
+ } else {
+ val grouped = groups.zipWithIndex.map { case (group, i) =>
+ splitCall(s"${funcName}_group${level}_$i",
+ group.flatMap(_.arguments).distinctBy(_._2).toSeq, returnType,
+ makeSplitFunction(foldFunctions(group.map(_.call).toSeq)))
+ }.toSeq
+ groupSplitCalls(grouped, funcName, returnType, makeSplitFunction,
foldFunctions, level + 1)
+ }
+ }
+ }
+
+ /** The parameter slots of a method taking `arguments`, `this` included. */
+ private def paramSlots(arguments: Seq[(String, String)]): Int =
+ 1 + arguments.map { case (javaType, _) =>
+ if (javaType == JAVA_LONG || javaType == JAVA_DOUBLE) 2 else 1
+ }.sum
+
+ /**
+ * The arguments of a function split out of whole stage generated code so
that the code generated
+ * for `sources` compiles inside it: the local variables that code reads
from the operator's
+ * method. None where the code cannot be split, because a local it reads
cannot be told from
+ * `sources` or cannot be passed.
+ *
+ * The code reads its inputs from `currentVars`, from the states of
subexpression elimination
+ * and, where the operator sets one, from `INPUT_ROW`; a field or a literal
among them is read as
+ * it stands and needs no argument. A state is taken as its own variables
without looking below
+ * it, the way `CodeGenerator.getLocalInputVariableValues` takes it, so that
the arguments are
+ * among those of the functions an aggregate or `ExpandExec` splits out of
its own code, which is
+ * where these calls then sit. A `CommonExpressionRef` computes its
definition where it is
+ * reached, so what the definition reads is read here too.
+ *
+ * It refuses, as `CommonExprSlots.methodArgs` does and for the reasons it
gives:
+ * - an input the operator has not evaluated yet, whose code was generated
into this code and
+ * names locals of the operator producing the row;
+ * - a value no parameter can carry: a `SimpleExprValue`, or a slot of a
compacted mutable
+ * state array such as `mutableStateArray_0[3]`, which is not a name;
+ * - more parameter slots than a method can take, `extraArguments` included.
+ * And two it has of its own:
+ * - a state for a node that overrides `genCode` and so may generate its
child instead of reading
+ * the state (`Alias`, `Collate` and a `Cast` that changes no structure
do), so that what the
+ * code reads below it is not in the state;
+ * - a `LambdaVariable`. Its value is a field that the loop of its object
expression assigns,
+ * which a function could read as it is, but no test reaches such a
split, so the code keeps
+ * its place.
+ * `INPUT_ROW` is passed only when the code reads it: through an input the
operator has not put
+ * in `currentVars`, or through a `CommonExpressionRef`, whose definition
`CommonExprSlots` may
+ * make a method that takes `INPUT_ROW` wherever it is set. An operator may
leave `INPUT_ROW`
+ * naming a local that is not in scope where other code is, as a join does
after generating the
+ * variables of its build side, which is why it is not passed otherwise.
+ */
+ private def wholeStageSplitArguments(
+ sources: Seq[Expression],
+ extraArguments: Seq[(String, String)]): Option[Seq[(String, String)]] = {
+ val args = mutable.LinkedHashMap.empty[String, VariableValue]
+ def isName(name: String): Boolean = name.nonEmpty &&
+ Character.isJavaIdentifierStart(name.head) &&
name.forall(Character.isJavaIdentifierPart)
+ def canPass(v: ExprValue): Boolean = v match {
+ case local: VariableValue =>
+ val name = local.variableName
+ // A field needs no argument, and passing one is refused by
`splitExpressions`.
+ if (isName(name) && !mutableStateNames.contains(name)) {
+ args.getOrElseUpdate(name, local)
+ }
+ isName(name)
+ case _: GlobalValue | _: LiteralValue => true
+ case _ => false
+ }
+ // `Expression.genCode` is what reads a state; a node that overrides it
may not.
+ def readsSubExprState(e: Expression): Boolean = e match {
+ case c: Cast => !c.generatesChildCode
+ case _ => e.getClass.getMethod("genCode",
classOf[CodegenContext]).getDeclaringClass ==
+ classOf[Expression]
+ }
+ var needsRow = false
+ var possible = sources.nonEmpty
+ val visited = mutable.HashSet.empty[Long]
+ // The definitions of a `With` in the trees walked here, which are in
`currentCommonExprs`
+ // only while that `With` is generated.
+ val nestedDefs = mutable.HashMap.empty[Long, Expression]
+ val toVisit = mutable.Stack[Expression](sources: _*)
+ while (possible && toVisit.nonEmpty) {
+ val next = toVisit.pop()
+ subExprEliminationExprs.get(ExpressionEquals(next)) match {
+ case Some(state) =>
+ possible = readsSubExprState(next) &&
+ canPass(state.eval.value) && canPass(state.eval.isNull)
+ case None =>
+ next match {
+ case ref: BoundReference if currentVars(ref.ordinal) != null =>
+ val input = currentVars(ref.ordinal)
+ possible = input.code == EmptyBlock && canPass(input.value) &&
canPass(input.isNull)
+ case _: BoundReference =>
+ needsRow = true
+ case _: LambdaVariable =>
+ possible = false
+ case w: With =>
+ w.defs.foreach(d => nestedDefs.put(d.id.id, d.child))
+ toVisit.push(w.child)
+ case ref: CommonExpressionRef =>
+ // The call to a definition made a method passes `INPUT_ROW`
wherever it is set, and
+ // whether the definition is made one depends on the length of
its code
+ // (`CommonExprSlots.build`), which is generated only after
these arguments are
+ // computed. So the row is taken wherever it is set: for a
definition that ends up
+ // inlined, an argument the method does not read. That is safe
where references
+ // reach code generation - an aggregate, where `INPUT_ROW` is in
scope -
+ // `RewriteWithExpression` having inlined a `With` in a join
condition.
+ needsRow |= INPUT_ROW != null
+ 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)
+ }
+ }
+ }
+ if (needsRow) {
+ possible &&= INPUT_ROW != null && canPass(JavaCode.variable(INPUT_ROW,
classOf[InternalRow]))
+ }
+ val params = args.values.toSeq
+ val extraSlots = extraArguments.map { case (javaType, _) =>
Review Comment:
`extraSlots` is exactly `paramSlots(extraArguments) - 1`, and both repeat
the rule in `calculateParamLength`/`calculateParamLengthFromExprValues`. How
about this, so that the block check here and the group check in
`groupSplitCalls` use the same function?
```scala
val arguments = params.map(v => typeName(v.javaType) -> v.variableName)
Option.when(possible && isValidParamLength(paramSlots(arguments ++
extraArguments)))(arguments)
```
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -1188,6 +1198,270 @@ class CodegenContext extends Logging {
}
}
+ /**
+ * `splitExpressionsWithCurrentInputs` for code whose sources 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, so `sources` gives those, one sequence for each entry of
`expressions`. Each function
+ * split out takes the locals its own entries read (see
`wholeStageSplitArguments`), and a
+ * function that groups calls takes what its calls take. The code stays in
one piece where the
+ * code of any function reads a local no function could take. Outside whole
stage codegen this is
+ * the method above.
+ *
+ * It is an overload rather than a new parameter of the method above, whose
signature code
+ * outside Spark is compiled against.
+ *
+ * @param expressions the codes to evaluate expressions.
+ * @param sources the expressions each entry of `expressions` was generated
from.
+ * @param funcName the split function name base.
+ * @param extraArguments the list of (type, name) of the arguments of the
split functions,
+ * except for 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.
+ */
+ def splitExpressionsWithCurrentInputs(
Review Comment:
This adds a public overload that takes `sources` as a sequence parallel to
`expressions`, checked only by length. It can't have default arguments because
the other overload has them, so callers need to spell out `extraArguments =
Nil` and the rest. Since catalyst is excluded from MiMa
(`org.apache.spark.sql.catalyst.*`), what about a `private[sql]` method with a
different name that takes `(code, sources)` pairs and defaults? That removes
the `require` and makes adoption easier for the other callers (`Coalesce`,
`In`, `Least`/`Greatest`, `CreateNamedStruct`, `ConcatWs`, ...), which still
never split under whole-stage codegen.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -1528,7 +1809,12 @@ class CodegenContext extends Logging {
// elimination.
val commonExprs = equivalentExpressions.getCommonSubexpressions
- val nonSplitCode = {
+ // The code of the non-split pass is discarded when the split pass below
is taken, and with it
Review Comment:
Only operators call `subexpressionEliminationForWholeStageCodegen`
(`ExpandExec`, `AggregateCodegenSupport`, `HashAggregateExec`, `ProjectExec`,
`FilterExec`). No expression calls it during another pass's `expr.genCode`, so
`outerDiscardable` is always null, and the merge at L1840-1842 is dead code.
Since all function names come from `freshName`, a local diff of the
`classFunctions` key sets before and after the non-split pass would give the
same set, without the context-wide `discardableFunctions` var, the hook in
`addNewFunctionToClass` or the try/finally.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -1188,6 +1198,270 @@ class CodegenContext extends Logging {
}
}
+ /**
+ * `splitExpressionsWithCurrentInputs` for code whose sources 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, so `sources` gives those, one sequence for each entry of
`expressions`. Each function
+ * split out takes the locals its own entries read (see
`wholeStageSplitArguments`), and a
+ * function that groups calls takes what its calls take. The code stays in
one piece where the
+ * code of any function reads a local no function could take. Outside whole
stage codegen this is
+ * the method above.
+ *
+ * It is an overload rather than a new parameter of the method above, whose
signature code
+ * outside Spark is compiled against.
+ *
+ * @param expressions the codes to evaluate expressions.
+ * @param sources the expressions each entry of `expressions` was generated
from.
+ * @param funcName the split function name base.
+ * @param extraArguments the list of (type, name) of the arguments of the
split functions,
+ * except for 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.
+ */
+ def splitExpressionsWithCurrentInputs(
+ expressions: Seq[String],
+ sources: Seq[Seq[Expression]],
+ funcName: String,
+ extraArguments: Seq[(String, String)],
+ returnType: String,
+ makeSplitFunction: String => String,
+ foldFunctions: Seq[String] => String): String = {
+ require(sources.length == expressions.length,
+ s"${sources.length} source sequences for ${expressions.length} pieces of
code")
+ if (currentVars == null) {
+ splitExpressionsWithCurrentInputs(
+ expressions, funcName, extraArguments, returnType, makeSplitFunction,
foldFunctions)
+ } else if (!SQLConf.get.wholeStageSplitExpressions) {
+ expressions.mkString("\n")
+ } else {
+ splitWholeStage(expressions, sources, funcName, extraArguments,
returnType,
+ makeSplitFunction, foldFunctions).getOrElse(expressions.mkString("\n"))
+ }
+ }
+
+ /**
+ * The whole stage split of `splitExpressionsWithCurrentInputs`: None where
the code is one
+ * block, or where the code of a block reads a local no function could take.
+ */
+ private def splitWholeStage(
+ expressions: Seq[String],
+ sources: Seq[Seq[Expression]],
+ funcName: String,
+ extraArguments: Seq[(String, String)],
+ returnType: String,
+ makeSplitFunction: String => String,
+ foldFunctions: Seq[String] => String): Option[String] = {
+ val blocks = buildCodeBlocksOf(expressions)
+ if (blocks.length <= 1) {
Review Comment:
The whole-stage split is decided only by `blocks.length > 1`, i.e. by
`methodSplitThreshold` (1024 chars of source), not by the 8000-byte limit the
split is meant to avoid. In the benchmark's shape, a branch with ANSI
arithmetic is ~740 chars, so even the 2-branch CASE WHEN already gets its ELSE
split into a separate method.
The PR's own results show the cost where the unsplit method is still
JIT-compilable. At 64 branches (6049 bytes unsplit): 116 vs 103 ms (JDK 17),
126 vs 104 (JDK 21), 120 vs 101 (JDK 25). The only ways back are the
all-or-nothing internal conf, which brings back the 14-28x cliff from 128
branches, or the global `methodSplitThreshold`, which also drives
`CommonExprSlots`, the subexpression-elimination split, `ExpandExec` and the
aggregate split.
Note that splitting less per expression would make the comment on
`groupSplitCalls` (several CASE WHENs in one method) worse, so this probably
needs a stage-level decision, e.g. based on the compiled `maxMethodCodeSize`.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/WholeStageCodegenSuite.scala:
##########
@@ -1799,4 +1799,256 @@ 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 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: the split methods take nullable, string and subexpression
inputs") {
+ // Read from a file, so the inputs are columns. `upper(s)` is in every
condition, so
+ // subexpression elimination computes it once into local variables, which
the split methods
+ // take as parameters beside the columns.
+ withTempPath { path =>
+ spark.range(200)
+ .selectExpr(
+ "if(id % 7 = 0, null, concat('v', cast(id % 50 as string))) AS s",
+ "if(id % 11 = 0, null, id) AS k")
+ .write.parquet(path.getCanonicalPath)
+ val branches = (0 until 50).map { i =>
+ s"WHEN upper(s) = 'V$i' THEN concat(s, cast(k + $i as string))"
+ }.mkString(" ")
+ def df: DataFrame =
+ spark.read.parquet(path.getCanonicalPath).selectExpr(s"CASE $branches
ELSE upper(s) END")
+ assert(splitsCaseWhen(df))
+ checkAnswer(df, withoutWholeStage(df))
+ }
+ }
+
+ test("SPARK-33301: each split method takes only what its own branches read")
{
+ // Every branch reads its own nullable `bigint` column, three parameter
slots each, so the
+ // union over all branches is past the 255 slots a method can take, while
each method takes
+ // the columns of the branches it holds, and the methods grouping the
calls the union of what
+ // their calls take. Each column is read twice, so the projection
evaluates it before the
+ // CASE WHEN.
+ withTempPath { path =>
+ val columns = 90
+ spark.range(100)
+ .selectExpr((1 to columns).map(i => s"if(id % $i = 0, null, id + $i)
AS c$i"): _*)
+ .write.parquet(path.getCanonicalPath)
+ val branches = (1 to columns).map(i => s"WHEN c$i > 150 THEN
c$i").mkString(" ")
+ def df: DataFrame =
+ spark.read.parquet(path.getCanonicalPath).selectExpr(s"CASE $branches
ELSE -1 END AS v")
+ assert(splitsCaseWhen(df))
+ checkAnswer(df, withoutWholeStage(df))
+ }
+ }
+
+ test("SPARK-33301: a CASE WHEN whose methods would take too many parameters
stays in one piece") {
+ // Every branch reads all three columns, and no expression repeats, so
that subexpression
+ // elimination, which the same limit governs, has nothing to split. Under
the default limit
+ // the same CASE WHEN splits, so it is the limit that keeps it in one
piece.
+ withTempPath { path =>
+ spark.range(50).selectExpr("id AS a", "id + 1 AS b", "id + 2 AS c")
+ .write.parquet(path.getCanonicalPath)
+ val branches = (1 to 100).map(k => s"WHEN a + $k = b THEN c *
$k").mkString(" ")
+ def df: DataFrame =
+ spark.read.parquet(path.getCanonicalPath).selectExpr(s"CASE $branches
ELSE 0 END AS v")
+ assert(splitsCaseWhen(df))
+ withSQLConf("spark.sql.CodeGenerator.validParamLength" -> "4") {
+ assert(!splitsCaseWhen(df))
+ checkAnswer(df, withoutWholeStage(df))
+ }
+ }
+ }
+
+ test("SPARK-33301: a CASE WHEN under a keyed aggregate takes its input, not
the buffer row") {
+ // The aggregate's update code reads the buffer from `INPUT_ROW` and the
input from local
+ // variables. The CASE WHEN reads only the input, so its split methods
take that and not the
+ // buffer row, which is passed only to code that reads it.
+ withTempView("t") {
+ spark.range(1000).selectExpr("id % 13 AS k", "id % 300 AS
v").createOrReplaceTempView("t")
+ val query = s"SELECT k, sum(${largeCaseWhen("v", 300)}) AS s FROM t
GROUP BY k"
+ val df = sql(query)
+ assert(df.queryExecution.executedPlan.exists {
+ case w: WholeStageCodegenExec =>
w.child.isInstanceOf[HashAggregateExec]
+ case _ => false
+ })
+ assert(splitsCaseWhen(df))
+ assert(!splitCaseWhenTakesRow(df))
+ checkAnswer(df, withoutWholeStage(sql(query)))
+ }
+ }
+
+ test("SPARK-33301: a common expression in a branch under a keyed aggregate
takes the buffer") {
+ // `nullif` is a `With`, which stays in the plan in a branch other than
the first, and whose
+ // definition, made a method, is called with `INPUT_ROW`: the aggregate's
buffer, which the
+ // split methods therefore take.
+ withTempView("t") {
+ spark.range(1000).selectExpr("id % 7 AS k", "CAST(id % 400 AS INT) AS v")
+ .createOrReplaceTempView("t")
+ val branches = "WHEN v = 0 THEN 0 WHEN v = 1 THEN nullif(nullif(v * 3,
1), 2) " +
+ (2 to 300).map(k => s"WHEN v = $k THEN v * $k").mkString(" ")
+ val query = s"SELECT k, sum(CASE $branches ELSE 0 END) AS s FROM t GROUP
BY k"
+ val df = sql(query)
+ assert(splitsCaseWhen(df))
+ assert(splitCaseWhenTakesRow(df))
+ checkAnswer(df, withoutWholeStage(sql(query)))
+ }
+ }
+
+ test("SPARK-33301: max_by under a keyed aggregate, whose update reads the
buffer") {
+ // `max_by` updates and merges its buffer with CASE WHENs of three
branches that read the
+ // buffer through `INPUT_ROW`, so their split methods take the row. The
value is a number: a
+ // string one is not updatable in place, and Spark plans a sort aggregate
outside the stage.
+ withTempView("t") {
+ spark.range(1000).selectExpr("id % 7 AS k", "id % 97 AS
o").createOrReplaceTempView("t")
+ val query = "SELECT k, max_by(o * 3, o), min_by(o * 3, o) FROM t GROUP
BY k"
+ assert(splitsCaseWhen(sql(query)))
+ assert(splitCaseWhenTakesRow(sql(query)))
+ checkAnswer(sql(query), withoutWholeStage(sql(query)))
+ }
+ }
+
+ test("SPARK-33301: a large CASE WHEN in a join condition compiles") {
+ // A join generates its condition inside the loop over the matches, where
`INPUT_ROW` names
+ // the row of the build side; the split methods take it only if the code
reads it.
+ withTempView("l", "r") {
+ spark.range(200).selectExpr("id % 20 AS k", "id % 300 AS
v").createOrReplaceTempView("l")
+ spark.range(100).selectExpr("id % 20 AS k", "id % 7 AS
w").createOrReplaceTempView("r")
+ val condition = s"${largeCaseWhen("l.v", 300)} + r.w > 50"
+ Seq(
+ s"SELECT /*+ MERGE(r) */ l.k, l.v, r.w FROM l JOIN r ON l.k = r.k AND
$condition",
+ s"SELECT /*+ BROADCAST(r) */ l.k, l.v, r.w FROM l JOIN r ON l.k = r.k
AND $condition",
+ s"SELECT /*+ MERGE(r) */ l.k, l.v, r.w FROM l LEFT JOIN r " +
+ s"ON l.k = r.k AND ${largeCaseWhen("l.v", 300)} > 50"
+ ).foreach { query =>
+ val df = sql(query)
+ assert(splitsCaseWhen(df), query)
+ checkAnswer(df, withoutWholeStage(sql(query)))
+ }
+ }
+ }
+
+ test("SPARK-33301: a CASE WHEN in an aggregate over a ROLLUP") {
+ // The aggregate reads the rows `ExpandExec` produces, through the
parameters of its own
+ // consume method.
+ withTempView("t") {
+ spark.range(300)
+ .selectExpr("id % 300 AS v", "concat('g', cast(id % 3 AS string)) AS
g")
+ .createOrReplaceTempView("t")
+ val query =
+ s"SELECT g, sum(${largeCaseWhen("v", 300)}) AS s FROM t GROUP BY
ROLLUP(g, v % 2)"
+ val df = sql(query)
+ assert(splitsCaseWhen(df))
+ checkAnswer(df, withoutWholeStage(sql(query)))
+ }
+ }
+
+ test("SPARK-33301: a CASE WHEN that is itself a split common subexpression")
{
+ // The same CASE WHEN twice is one subexpression, split into a method of
its own; the CASE
+ // WHEN's split methods are called from inside it.
+ val caseWhen = largeCaseWhen("id", 300)
+ val df = spark.range(10).selectExpr(s"($caseWhen) + 1 AS a", s"($caseWhen)
+ 2 AS b")
+ assert(splitsCaseWhen(df))
+ checkAnswer(df, (0L until 10L).map(i => Row(i * i + 1, i * i + 2)))
+ }
+
+ test("SPARK-33301: a split common subexpression under a keyed aggregate
takes the buffer") {
+ // The same CASE WHEN in two aggregates is one subexpression, split into a
`subExpr` method
+ // of its own. A `nullif` branch makes the CASE WHEN's split methods take
the aggregate's
+ // buffer, so the `subExpr` method, which calls them, must take it too.
+ withTempView("t") {
+ spark.range(1000).selectExpr("id % 7 AS k", "CAST(id % 400 AS INT) AS v")
+ .createOrReplaceTempView("t")
+ val branches = "WHEN v = 0 THEN 0 WHEN v = 1 THEN nullif(nullif(v * 3,
1), 2) " +
+ (2 to 300).map(k => s"WHEN v = $k THEN v * $k").mkString(" ")
+ val caseWhen = s"CASE $branches ELSE 0 END"
+ val query = s"SELECT k, sum($caseWhen), max($caseWhen) FROM t GROUP BY k"
+ val df = sql(query)
+ assert(splitCaseWhenTakesRow(df))
+ val subExpr = "(?s)private void \\w*subExpr_\\d+\\([^)]*InternalRow
[^{]*\\{" +
+ "(?:(?!\\nprivate ).)*caseWhen_\\d"
+ assert(genCode(df).exists(c => subExpr.r.findFirstIn(c.body).nonEmpty))
+ checkAnswer(df, withoutWholeStage(sql(query)))
+ }
+ }
+
+ test("SPARK-33301: functions added in a discarded pass of subexpression
elimination go") {
Review Comment:
This test checks only `splitsCaseWhen` and `checkAnswer`, so it would still
pass if the removal at CodeGenerator.scala L1857-1859 were deleted entirely.
The first pass's CASE WHEN, `computeCommonExpr` and group methods would just
remain as unused private methods, which compile fine. The test only catches a
partial removal. Could we also assert that the functions are actually gone,
e.g. by counting the `caseWhen_` methods or checking that the first pass's
function names are absent from the generated code?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -1188,6 +1198,270 @@ class CodegenContext extends Logging {
}
}
+ /**
+ * `splitExpressionsWithCurrentInputs` for code whose sources 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, so `sources` gives those, one sequence for each entry of
`expressions`. Each function
+ * split out takes the locals its own entries read (see
`wholeStageSplitArguments`), and a
+ * function that groups calls takes what its calls take. The code stays in
one piece where the
+ * code of any function reads a local no function could take. Outside whole
stage codegen this is
+ * the method above.
+ *
+ * It is an overload rather than a new parameter of the method above, whose
signature code
+ * outside Spark is compiled against.
+ *
+ * @param expressions the codes to evaluate expressions.
+ * @param sources the expressions each entry of `expressions` was generated
from.
+ * @param funcName the split function name base.
+ * @param extraArguments the list of (type, name) of the arguments of the
split functions,
+ * except for 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.
+ */
+ def splitExpressionsWithCurrentInputs(
+ expressions: Seq[String],
+ sources: Seq[Seq[Expression]],
+ funcName: String,
+ extraArguments: Seq[(String, String)],
+ returnType: String,
+ makeSplitFunction: String => String,
+ foldFunctions: Seq[String] => String): String = {
+ require(sources.length == expressions.length,
+ s"${sources.length} source sequences for ${expressions.length} pieces of
code")
+ if (currentVars == null) {
+ splitExpressionsWithCurrentInputs(
+ expressions, funcName, extraArguments, returnType, makeSplitFunction,
foldFunctions)
+ } else if (!SQLConf.get.wholeStageSplitExpressions) {
+ expressions.mkString("\n")
+ } else {
+ splitWholeStage(expressions, sources, funcName, extraArguments,
returnType,
+ makeSplitFunction, foldFunctions).getOrElse(expressions.mkString("\n"))
+ }
+ }
+
+ /**
+ * The whole stage split of `splitExpressionsWithCurrentInputs`: None where
the code is one
+ * block, or where the code of a block reads a local no function could take.
+ */
+ private def splitWholeStage(
+ expressions: Seq[String],
+ sources: Seq[Seq[Expression]],
+ funcName: String,
+ extraArguments: Seq[(String, String)],
+ returnType: String,
+ makeSplitFunction: String => String,
+ foldFunctions: Seq[String] => String): Option[String] = {
+ val blocks = buildCodeBlocksOf(expressions)
+ if (blocks.length <= 1) {
+ None
+ } else {
+ val blockArguments = blocks.map { case (_, entries) =>
+ wholeStageSplitArguments(entries.flatMap(sources), extraArguments)
+ }
+ Option.when(blockArguments.forall(_.isDefined)) {
+ val func = freshName(funcName)
+ val calls = blocks.zip(blockArguments.flatten).zipWithIndex.map {
+ case (((body, _), arguments), i) =>
+ splitCall(s"${func}_$i", arguments ++ extraArguments, returnType,
+ makeSplitFunction(body))
+ }
+ foldFunctions(
+ groupSplitCalls(calls, func, returnType, makeSplitFunction,
foldFunctions).map(_.call))
+ }
+ }
+ }
+
+ /** A call to a function split out of whole stage generated code, and the
arguments it takes. */
+ private case class SplitCall(call: String, arguments: Seq[(String, String)])
+
+ /** Adds the function `name` taking `arguments` and running `body`, and
returns a call to it. */
+ private def splitCall(
+ name: String,
+ arguments: Seq[(String, String)],
+ returnType: String,
+ body: String): SplitCall = {
+ if (Utils.isTesting) {
+ // Passing global variables to the split method is dangerous, as any
mutating to it is
+ // ignored and may lead to unexpected behavior.
+ arguments.foreach { case (_, argument) =>
+ assert(!mutableStateNames.contains(argument),
+ s"split function argument $argument cannot be a global variable.")
+ }
+ }
+ val argString = arguments.map { case (t, argument) => s"$t $argument"
}.mkString(", ")
+ val code =
+ s"""
+ |private $returnType $name($argString) {
+ | $body
+ |}
+ """.stripMargin
+ SplitCall(s"${addNewFunction(name,
code)}(${arguments.map(_._2).mkString(", ")})", arguments)
+ }
+
+ /**
+ * `groupFunctionCalls` for calls that take different arguments: a group
takes the union of what
+ * its calls take, and is closed when one more call would take it past
+ * `SPLIT_CALLS_PER_GROUP` calls or past the parameters a method can have.
Where no group can
+ * hold two calls, the calls are left as they are.
+ */
+ private def groupSplitCalls(
Review Comment:
`groupSplitCalls` and `splitCall` duplicate `groupFunctionCalls` and the
function template and `Utils.isTesting` assertion of `splitExpressions`
(SPARK-59783), with per-call arguments. With uniform arguments,
`groupSplitCalls(calls.map(SplitCall(_, args)), ...)` gives exactly what
`groupFunctionCalls` gives, so could `splitExpressions` reuse it? Otherwise the
grouping policy (`SPLIT_CALLS_PER_GROUP`, the `splitCallsPerMethod` budget, the
`_group${level}_$i` naming) exists twice, and the whole-stage and
non-whole-stage paths can drift silently. The no-progress guard is already only
here.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -1188,6 +1198,270 @@ class CodegenContext extends Logging {
}
}
+ /**
+ * `splitExpressionsWithCurrentInputs` for code whose sources 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, so `sources` gives those, one sequence for each entry of
`expressions`. Each function
+ * split out takes the locals its own entries read (see
`wholeStageSplitArguments`), and a
+ * function that groups calls takes what its calls take. The code stays in
one piece where the
+ * code of any function reads a local no function could take. Outside whole
stage codegen this is
+ * the method above.
+ *
+ * It is an overload rather than a new parameter of the method above, whose
signature code
+ * outside Spark is compiled against.
+ *
+ * @param expressions the codes to evaluate expressions.
+ * @param sources the expressions each entry of `expressions` was generated
from.
+ * @param funcName the split function name base.
+ * @param extraArguments the list of (type, name) of the arguments of the
split functions,
+ * except for 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.
+ */
+ def splitExpressionsWithCurrentInputs(
+ expressions: Seq[String],
+ sources: Seq[Seq[Expression]],
+ funcName: String,
+ extraArguments: Seq[(String, String)],
+ returnType: String,
+ makeSplitFunction: String => String,
+ foldFunctions: Seq[String] => String): String = {
+ require(sources.length == expressions.length,
+ s"${sources.length} source sequences for ${expressions.length} pieces of
code")
+ if (currentVars == null) {
+ splitExpressionsWithCurrentInputs(
+ expressions, funcName, extraArguments, returnType, makeSplitFunction,
foldFunctions)
+ } else if (!SQLConf.get.wholeStageSplitExpressions) {
+ expressions.mkString("\n")
+ } else {
+ splitWholeStage(expressions, sources, funcName, extraArguments,
returnType,
+ makeSplitFunction, foldFunctions).getOrElse(expressions.mkString("\n"))
+ }
+ }
+
+ /**
+ * The whole stage split of `splitExpressionsWithCurrentInputs`: None where
the code is one
+ * block, or where the code of a block reads a local no function could take.
+ */
+ private def splitWholeStage(
+ expressions: Seq[String],
+ sources: Seq[Seq[Expression]],
+ funcName: String,
+ extraArguments: Seq[(String, String)],
+ returnType: String,
+ makeSplitFunction: String => String,
+ foldFunctions: Seq[String] => String): Option[String] = {
+ val blocks = buildCodeBlocksOf(expressions)
+ if (blocks.length <= 1) {
+ None
+ } else {
+ val blockArguments = blocks.map { case (_, entries) =>
+ wholeStageSplitArguments(entries.flatMap(sources), extraArguments)
Review Comment:
`entries.flatMap(sources)` indexes `sources` by position, but
`CaseWhen.branches` is a `List` after `withNewChildrenInternal`
(`Iterator.toSeq` in Scala 2.13). So this is O(B^2) in the number of branches:
~4.5M list hops at 3000 branches, twice when the CASE WHEN is a common
subexpression. It's minor, but `sources.toIndexedSeq`, or passing the sources
along with the blocks, fixes it. Stopping at the first `None` would also avoid
walking the remaining blocks after a refusal.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -1188,6 +1198,270 @@ class CodegenContext extends Logging {
}
}
+ /**
+ * `splitExpressionsWithCurrentInputs` for code whose sources 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, so `sources` gives those, one sequence for each entry of
`expressions`. Each function
+ * split out takes the locals its own entries read (see
`wholeStageSplitArguments`), and a
+ * function that groups calls takes what its calls take. The code stays in
one piece where the
+ * code of any function reads a local no function could take. Outside whole
stage codegen this is
+ * the method above.
+ *
+ * It is an overload rather than a new parameter of the method above, whose
signature code
+ * outside Spark is compiled against.
+ *
+ * @param expressions the codes to evaluate expressions.
+ * @param sources the expressions each entry of `expressions` was generated
from.
+ * @param funcName the split function name base.
+ * @param extraArguments the list of (type, name) of the arguments of the
split functions,
+ * except for 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.
+ */
+ def splitExpressionsWithCurrentInputs(
+ expressions: Seq[String],
+ sources: Seq[Seq[Expression]],
+ funcName: String,
+ extraArguments: Seq[(String, String)],
+ returnType: String,
+ makeSplitFunction: String => String,
+ foldFunctions: Seq[String] => String): String = {
+ require(sources.length == expressions.length,
+ s"${sources.length} source sequences for ${expressions.length} pieces of
code")
+ if (currentVars == null) {
+ splitExpressionsWithCurrentInputs(
+ expressions, funcName, extraArguments, returnType, makeSplitFunction,
foldFunctions)
+ } else if (!SQLConf.get.wholeStageSplitExpressions) {
+ expressions.mkString("\n")
+ } else {
+ splitWholeStage(expressions, sources, funcName, extraArguments,
returnType,
+ makeSplitFunction, foldFunctions).getOrElse(expressions.mkString("\n"))
+ }
+ }
+
+ /**
+ * The whole stage split of `splitExpressionsWithCurrentInputs`: None where
the code is one
+ * block, or where the code of a block reads a local no function could take.
+ */
+ private def splitWholeStage(
+ expressions: Seq[String],
+ sources: Seq[Seq[Expression]],
+ funcName: String,
+ extraArguments: Seq[(String, String)],
+ returnType: String,
+ makeSplitFunction: String => String,
+ foldFunctions: Seq[String] => String): Option[String] = {
+ val blocks = buildCodeBlocksOf(expressions)
+ if (blocks.length <= 1) {
+ None
+ } else {
+ val blockArguments = blocks.map { case (_, entries) =>
+ wholeStageSplitArguments(entries.flatMap(sources), extraArguments)
+ }
+ Option.when(blockArguments.forall(_.isDefined)) {
+ val func = freshName(funcName)
+ val calls = blocks.zip(blockArguments.flatten).zipWithIndex.map {
+ case (((body, _), arguments), i) =>
+ splitCall(s"${func}_$i", arguments ++ extraArguments, returnType,
+ makeSplitFunction(body))
+ }
+ foldFunctions(
+ groupSplitCalls(calls, func, returnType, makeSplitFunction,
foldFunctions).map(_.call))
+ }
+ }
+ }
+
+ /** A call to a function split out of whole stage generated code, and the
arguments it takes. */
+ private case class SplitCall(call: String, arguments: Seq[(String, String)])
+
+ /** Adds the function `name` taking `arguments` and running `body`, and
returns a call to it. */
+ private def splitCall(
+ name: String,
+ arguments: Seq[(String, String)],
+ returnType: String,
+ body: String): SplitCall = {
+ if (Utils.isTesting) {
+ // Passing global variables to the split method is dangerous, as any
mutating to it is
+ // ignored and may lead to unexpected behavior.
+ arguments.foreach { case (_, argument) =>
+ assert(!mutableStateNames.contains(argument),
+ s"split function argument $argument cannot be a global variable.")
+ }
+ }
+ val argString = arguments.map { case (t, argument) => s"$t $argument"
}.mkString(", ")
+ val code =
+ s"""
+ |private $returnType $name($argString) {
+ | $body
+ |}
+ """.stripMargin
+ SplitCall(s"${addNewFunction(name,
code)}(${arguments.map(_._2).mkString(", ")})", arguments)
+ }
+
+ /**
+ * `groupFunctionCalls` for calls that take different arguments: a group
takes the union of what
+ * its calls take, and is closed when one more call would take it past
+ * `SPLIT_CALLS_PER_GROUP` calls or past the parameters a method can have.
Where no group can
+ * hold two calls, the calls are left as they are.
+ */
+ private def groupSplitCalls(
+ calls: Seq[SplitCall],
+ funcName: String,
+ returnType: String,
+ makeSplitFunction: String => String,
+ foldFunctions: Seq[String] => String,
+ level: Int = 0): Seq[SplitCall] = {
+ if (calls.isEmpty ||
+ calls.length <=
splitCallsPerMethod(calls.maxBy(_.arguments.length).arguments)) {
+ calls
+ } else {
+ val groups = mutable.ArrayBuffer.empty[mutable.ArrayBuffer[SplitCall]]
+ var taken = Seq.empty[(String, String)]
+ calls.foreach { c =>
+ val union = (taken ++ c.arguments).distinctBy(_._2)
+ if (groups.isEmpty || groups.last.length == SPLIT_CALLS_PER_GROUP ||
+ !isValidParamLength(paramSlots(union))) {
+ groups += mutable.ArrayBuffer(c)
+ taken = c.arguments
+ } else {
+ groups.last += c
+ taken = union
+ }
+ }
+ if (groups.length == calls.length) {
+ calls
+ } else {
+ val grouped = groups.zipWithIndex.map { case (group, i) =>
+ splitCall(s"${funcName}_group${level}_$i",
+ group.flatMap(_.arguments).distinctBy(_._2).toSeq, returnType,
+ makeSplitFunction(foldFunctions(group.map(_.call).toSeq)))
+ }.toSeq
+ groupSplitCalls(grouped, funcName, returnType, makeSplitFunction,
foldFunctions, level + 1)
+ }
+ }
+ }
+
+ /** The parameter slots of a method taking `arguments`, `this` included. */
+ private def paramSlots(arguments: Seq[(String, String)]): Int =
+ 1 + arguments.map { case (javaType, _) =>
+ if (javaType == JAVA_LONG || javaType == JAVA_DOUBLE) 2 else 1
+ }.sum
+
+ /**
+ * The arguments of a function split out of whole stage generated code so
that the code generated
+ * for `sources` compiles inside it: the local variables that code reads
from the operator's
+ * method. None where the code cannot be split, because a local it reads
cannot be told from
+ * `sources` or cannot be passed.
+ *
+ * The code reads its inputs from `currentVars`, from the states of
subexpression elimination
+ * and, where the operator sets one, from `INPUT_ROW`; a field or a literal
among them is read as
+ * it stands and needs no argument. A state is taken as its own variables
without looking below
+ * it, the way `CodeGenerator.getLocalInputVariableValues` takes it, so that
the arguments are
+ * among those of the functions an aggregate or `ExpandExec` splits out of
its own code, which is
+ * where these calls then sit. A `CommonExpressionRef` computes its
definition where it is
+ * reached, so what the definition reads is read here too.
+ *
+ * It refuses, as `CommonExprSlots.methodArgs` does and for the reasons it
gives:
+ * - an input the operator has not evaluated yet, whose code was generated
into this code and
+ * names locals of the operator producing the row;
+ * - a value no parameter can carry: a `SimpleExprValue`, or a slot of a
compacted mutable
+ * state array such as `mutableStateArray_0[3]`, which is not a name;
+ * - more parameter slots than a method can take, `extraArguments` included.
+ * And two it has of its own:
+ * - a state for a node that overrides `genCode` and so may generate its
child instead of reading
+ * the state (`Alias`, `Collate` and a `Cast` that changes no structure
do), so that what the
+ * code reads below it is not in the state;
+ * - a `LambdaVariable`. Its value is a field that the loop of its object
expression assigns,
+ * which a function could read as it is, but no test reaches such a
split, so the code keeps
+ * its place.
+ * `INPUT_ROW` is passed only when the code reads it: through an input the
operator has not put
+ * in `currentVars`, or through a `CommonExpressionRef`, whose definition
`CommonExprSlots` may
+ * make a method that takes `INPUT_ROW` wherever it is set. An operator may
leave `INPUT_ROW`
+ * naming a local that is not in scope where other code is, as a join does
after generating the
+ * variables of its build side, which is why it is not passed otherwise.
+ */
+ private def wholeStageSplitArguments(
+ sources: Seq[Expression],
+ extraArguments: Seq[(String, String)]): Option[Seq[(String, String)]] = {
+ val args = mutable.LinkedHashMap.empty[String, VariableValue]
+ def isName(name: String): Boolean = name.nonEmpty &&
+ Character.isJavaIdentifierStart(name.head) &&
name.forall(Character.isJavaIdentifierPart)
+ def canPass(v: ExprValue): Boolean = v match {
+ case local: VariableValue =>
+ val name = local.variableName
+ // A field needs no argument, and passing one is refused by
`splitExpressions`.
+ if (isName(name) && !mutableStateNames.contains(name)) {
+ args.getOrElseUpdate(name, local)
+ }
+ isName(name)
+ case _: GlobalValue | _: LiteralValue => true
+ case _ => false
+ }
+ // `Expression.genCode` is what reads a state; a node that overrides it
may not.
+ def readsSubExprState(e: Expression): Boolean = e match {
+ case c: Cast => !c.generatesChildCode
+ case _ => e.getClass.getMethod("genCode",
classOf[CodegenContext]).getDeclaringClass ==
+ classOf[Expression]
+ }
+ var needsRow = false
+ var possible = sources.nonEmpty
+ val visited = mutable.HashSet.empty[Long]
+ // The definitions of a `With` in the trees walked here, which are in
`currentCommonExprs`
+ // only while that `With` is generated.
+ val nestedDefs = mutable.HashMap.empty[Long, Expression]
+ val toVisit = mutable.Stack[Expression](sources: _*)
+ while (possible && toVisit.nonEmpty) {
+ val next = toVisit.pop()
+ subExprEliminationExprs.get(ExpressionEquals(next)) match {
+ case Some(state) =>
+ possible = readsSubExprState(next) &&
+ canPass(state.eval.value) && canPass(state.eval.isNull)
+ case None =>
+ next match {
+ case ref: BoundReference if currentVars(ref.ordinal) != null =>
+ val input = currentVars(ref.ordinal)
+ possible = input.code == EmptyBlock && canPass(input.value) &&
canPass(input.isNull)
+ case _: BoundReference =>
+ needsRow = true
+ case _: LambdaVariable =>
Review Comment:
Is this refusal needed? As the doc says, the value of
`objects.LambdaVariable` is a field that a split method can read as it is:
`prepareLambdaVariable` declares it with `forceInline = true`, so it's in
`mutableStateNames`. Without this arm, the leaf adds no argument anyway.
`NamedLambdaVariable` of higher-order functions is already handled that way
here, so this arm only adds another policy difference and the `objects` import.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/WholeStageCodegenSuite.scala:
##########
@@ -1799,4 +1799,256 @@ 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 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: the split methods take nullable, string and subexpression
inputs") {
+ // Read from a file, so the inputs are columns. `upper(s)` is in every
condition, so
+ // subexpression elimination computes it once into local variables, which
the split methods
+ // take as parameters beside the columns.
+ withTempPath { path =>
+ spark.range(200)
+ .selectExpr(
+ "if(id % 7 = 0, null, concat('v', cast(id % 50 as string))) AS s",
+ "if(id % 11 = 0, null, id) AS k")
+ .write.parquet(path.getCanonicalPath)
+ val branches = (0 until 50).map { i =>
+ s"WHEN upper(s) = 'V$i' THEN concat(s, cast(k + $i as string))"
+ }.mkString(" ")
+ def df: DataFrame =
+ spark.read.parquet(path.getCanonicalPath).selectExpr(s"CASE $branches
ELSE upper(s) END")
+ assert(splitsCaseWhen(df))
+ checkAnswer(df, withoutWholeStage(df))
+ }
+ }
+
+ test("SPARK-33301: each split method takes only what its own branches read")
{
+ // Every branch reads its own nullable `bigint` column, three parameter
slots each, so the
+ // union over all branches is past the 255 slots a method can take, while
each method takes
+ // the columns of the branches it holds, and the methods grouping the
calls the union of what
+ // their calls take. Each column is read twice, so the projection
evaluates it before the
+ // CASE WHEN.
+ withTempPath { path =>
+ val columns = 90
+ spark.range(100)
+ .selectExpr((1 to columns).map(i => s"if(id % $i = 0, null, id + $i)
AS c$i"): _*)
+ .write.parquet(path.getCanonicalPath)
+ val branches = (1 to columns).map(i => s"WHEN c$i > 150 THEN
c$i").mkString(" ")
+ def df: DataFrame =
+ spark.read.parquet(path.getCanonicalPath).selectExpr(s"CASE $branches
ELSE -1 END AS v")
+ assert(splitsCaseWhen(df))
+ checkAnswer(df, withoutWholeStage(df))
+ }
+ }
+
+ test("SPARK-33301: a CASE WHEN whose methods would take too many parameters
stays in one piece") {
+ // Every branch reads all three columns, and no expression repeats, so
that subexpression
+ // elimination, which the same limit governs, has nothing to split. Under
the default limit
+ // the same CASE WHEN splits, so it is the limit that keeps it in one
piece.
+ withTempPath { path =>
+ spark.range(50).selectExpr("id AS a", "id + 1 AS b", "id + 2 AS c")
+ .write.parquet(path.getCanonicalPath)
+ val branches = (1 to 100).map(k => s"WHEN a + $k = b THEN c *
$k").mkString(" ")
+ def df: DataFrame =
+ spark.read.parquet(path.getCanonicalPath).selectExpr(s"CASE $branches
ELSE 0 END AS v")
+ assert(splitsCaseWhen(df))
+ withSQLConf("spark.sql.CodeGenerator.validParamLength" -> "4") {
+ assert(!splitsCaseWhen(df))
+ checkAnswer(df, withoutWholeStage(df))
+ }
+ }
+ }
+
+ test("SPARK-33301: a CASE WHEN under a keyed aggregate takes its input, not
the buffer row") {
+ // The aggregate's update code reads the buffer from `INPUT_ROW` and the
input from local
+ // variables. The CASE WHEN reads only the input, so its split methods
take that and not the
+ // buffer row, which is passed only to code that reads it.
+ withTempView("t") {
+ spark.range(1000).selectExpr("id % 13 AS k", "id % 300 AS
v").createOrReplaceTempView("t")
+ val query = s"SELECT k, sum(${largeCaseWhen("v", 300)}) AS s FROM t
GROUP BY k"
+ val df = sql(query)
+ assert(df.queryExecution.executedPlan.exists {
+ case w: WholeStageCodegenExec =>
w.child.isInstanceOf[HashAggregateExec]
+ case _ => false
+ })
+ assert(splitsCaseWhen(df))
+ assert(!splitCaseWhenTakesRow(df))
+ checkAnswer(df, withoutWholeStage(sql(query)))
+ }
+ }
+
+ test("SPARK-33301: a common expression in a branch under a keyed aggregate
takes the buffer") {
+ // `nullif` is a `With`, which stays in the plan in a branch other than
the first, and whose
+ // definition, made a method, is called with `INPUT_ROW`: the aggregate's
buffer, which the
+ // split methods therefore take.
+ withTempView("t") {
+ spark.range(1000).selectExpr("id % 7 AS k", "CAST(id % 400 AS INT) AS v")
+ .createOrReplaceTempView("t")
+ val branches = "WHEN v = 0 THEN 0 WHEN v = 1 THEN nullif(nullif(v * 3,
1), 2) " +
+ (2 to 300).map(k => s"WHEN v = $k THEN v * $k").mkString(" ")
+ val query = s"SELECT k, sum(CASE $branches ELSE 0 END) AS s FROM t GROUP
BY k"
+ val df = sql(query)
+ assert(splitsCaseWhen(df))
+ assert(splitCaseWhenTakesRow(df))
+ checkAnswer(df, withoutWholeStage(sql(query)))
+ }
+ }
+
+ test("SPARK-33301: max_by under a keyed aggregate, whose update reads the
buffer") {
+ // `max_by` updates and merges its buffer with CASE WHENs of three
branches that read the
+ // buffer through `INPUT_ROW`, so their split methods take the row. The
value is a number: a
+ // string one is not updatable in place, and Spark plans a sort aggregate
outside the stage.
+ withTempView("t") {
+ spark.range(1000).selectExpr("id % 7 AS k", "id % 97 AS
o").createOrReplaceTempView("t")
+ val query = "SELECT k, max_by(o * 3, o), min_by(o * 3, o) FROM t GROUP
BY k"
+ assert(splitsCaseWhen(sql(query)))
+ assert(splitCaseWhenTakesRow(sql(query)))
+ checkAnswer(sql(query), withoutWholeStage(sql(query)))
+ }
+ }
+
+ test("SPARK-33301: a large CASE WHEN in a join condition compiles") {
+ // A join generates its condition inside the loop over the matches, where
`INPUT_ROW` names
+ // the row of the build side; the split methods take it only if the code
reads it.
+ withTempView("l", "r") {
+ spark.range(200).selectExpr("id % 20 AS k", "id % 300 AS
v").createOrReplaceTempView("l")
+ spark.range(100).selectExpr("id % 20 AS k", "id % 7 AS
w").createOrReplaceTempView("r")
+ val condition = s"${largeCaseWhen("l.v", 300)} + r.w > 50"
+ Seq(
+ s"SELECT /*+ MERGE(r) */ l.k, l.v, r.w FROM l JOIN r ON l.k = r.k AND
$condition",
+ s"SELECT /*+ BROADCAST(r) */ l.k, l.v, r.w FROM l JOIN r ON l.k = r.k
AND $condition",
+ s"SELECT /*+ MERGE(r) */ l.k, l.v, r.w FROM l LEFT JOIN r " +
+ s"ON l.k = r.k AND ${largeCaseWhen("l.v", 300)} > 50"
+ ).foreach { query =>
+ val df = sql(query)
+ assert(splitsCaseWhen(df), query)
+ checkAnswer(df, withoutWholeStage(sql(query)))
+ }
+ }
+ }
+
+ test("SPARK-33301: a CASE WHEN in an aggregate over a ROLLUP") {
+ // The aggregate reads the rows `ExpandExec` produces, through the
parameters of its own
+ // consume method.
+ withTempView("t") {
+ spark.range(300)
+ .selectExpr("id % 300 AS v", "concat('g', cast(id % 3 AS string)) AS
g")
+ .createOrReplaceTempView("t")
+ val query =
+ s"SELECT g, sum(${largeCaseWhen("v", 300)}) AS s FROM t GROUP BY
ROLLUP(g, v % 2)"
+ val df = sql(query)
+ assert(splitsCaseWhen(df))
+ checkAnswer(df, withoutWholeStage(sql(query)))
+ }
+ }
+
+ test("SPARK-33301: a CASE WHEN that is itself a split common subexpression")
{
+ // The same CASE WHEN twice is one subexpression, split into a method of
its own; the CASE
+ // WHEN's split methods are called from inside it.
+ val caseWhen = largeCaseWhen("id", 300)
+ val df = spark.range(10).selectExpr(s"($caseWhen) + 1 AS a", s"($caseWhen)
+ 2 AS b")
+ assert(splitsCaseWhen(df))
+ checkAnswer(df, (0L until 10L).map(i => Row(i * i + 1, i * i + 2)))
+ }
+
+ test("SPARK-33301: a split common subexpression under a keyed aggregate
takes the buffer") {
+ // The same CASE WHEN in two aggregates is one subexpression, split into a
`subExpr` method
+ // of its own. A `nullif` branch makes the CASE WHEN's split methods take
the aggregate's
+ // buffer, so the `subExpr` method, which calls them, must take it too.
+ withTempView("t") {
+ spark.range(1000).selectExpr("id % 7 AS k", "CAST(id % 400 AS INT) AS v")
+ .createOrReplaceTempView("t")
+ val branches = "WHEN v = 0 THEN 0 WHEN v = 1 THEN nullif(nullif(v * 3,
1), 2) " +
+ (2 to 300).map(k => s"WHEN v = $k THEN v * $k").mkString(" ")
+ val caseWhen = s"CASE $branches ELSE 0 END"
+ val query = s"SELECT k, sum($caseWhen), max($caseWhen) FROM t GROUP BY k"
+ val df = sql(query)
+ assert(splitCaseWhenTakesRow(df))
+ val subExpr = "(?s)private void \\w*subExpr_\\d+\\([^)]*InternalRow
[^{]*\\{" +
+ "(?:(?!\\nprivate ).)*caseWhen_\\d"
+ assert(genCode(df).exists(c => subExpr.r.findFirstIn(c.body).nonEmpty))
+ checkAnswer(df, withoutWholeStage(sql(query)))
+ }
+ }
+
+ test("SPARK-33301: functions added in a discarded pass of subexpression
elimination go") {
+ // Subexpression elimination generates its subexpressions once without
splitting and, when
+ // that code is too large, again split, discarding the first; the methods
the first pass
+ // added go with it, `computeCommonExpr` methods of a `nullif` included,
whose calls to the
+ // first pass's CASE WHEN methods would otherwise name methods no longer
declared.
+ val caseWhen = largeCaseWhen("id", 300)
+ val second = (1 to 300).map(k => s"WHEN id = $k THEN id + $k")
+ .mkString("CASE ", " ", " ELSE 0 END")
+ def df: DataFrame = spark.range(10).selectExpr(s"($caseWhen) + 1 AS a",
+ s"($caseWhen) + 2 AS b", s"IF(id < 0, 0, nullif($second, 0)) AS c",
+ s"IF(id < 0, 0, nullif($second, 0)) AS d")
+ assert(splitsCaseWhen(df))
+ checkAnswer(df, withoutWholeStage(df))
+ }
+
+ test("SPARK-33301: a CASE WHEN reading an input not evaluated yet stays in
one piece") {
+ // A projection defers reading a column it uses once to where it is used,
so the code reading
+ // each `c<i>` from the scan is inside the CASE WHEN's own code, and it
reads locals of the
+ // scan that no method could take.
+ withTempPath { path =>
+ spark.range(10).selectExpr("id % 40 AS k" +: (1 to 40).map(i => s"id +
$i AS c$i"): _*)
+ .write.parquet(path.getCanonicalPath)
+ val branches = (1 to 40).map(i => s"WHEN k = $i THEN c$i *
2").mkString(" ")
+ val df = spark.read.parquet(path.getCanonicalPath)
+ .selectExpr(s"CASE $branches ELSE -1 END AS v")
+ assert(!splitsCaseWhen(df))
Review Comment:
This asserts only `!splitsCaseWhen(df)`, so it would also pass if the split
were never attempted for another reason, e.g. if the code became one block.
Could we add a positive control? For example, the same CASE WHEN with each
`c<i>` read twice, so that `ProjectExec` evaluates them up front, and assert
that one splits.
--
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]
