MaxGekk commented on code in PR #59069:
URL: https://github.com/apache/spark/pull/59069#discussion_r4115802995
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -1188,6 +1198,263 @@ 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.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.
+ needsRow |= INPUT_ROW != null
Review Comment:
Thanks for trying to break it. Gating isn't available at that point: whether
a definition becomes a method depends on the length of its code
(`CommonExprSlots.build`), and that code is generated only after the split
arguments are computed. So I kept the conservative choice and added a comment
on why it's safe where references reach codegen. That's an aggregate, where
`INPUT_ROW` is in scope; `RewriteWithExpression` inlines a `With` in a join
condition, as you found. Done in 1d4975551c2.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -1188,6 +1198,263 @@ 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.length <=
splitCallsPerMethod(calls.maxBy(_.arguments.length).arguments)) {
Review Comment:
Done in 1d4975551c2. `groupSplitCalls` returns an empty `calls` as it is,
before `maxBy`.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/WholeStageCodegenSuite.scala:
##########
@@ -1799,4 +1799,232 @@ 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") {
Review Comment:
Good catch. The second note found more than a missing assert: the `max_by`
test never reached whole-stage codegen. A string value isn't updatable in
place, so Spark planned a sort aggregate outside the stage, and its CASE WHENs
didn't split even with `methodSplitThreshold` at 1. With a numeric value the
aggregate stays a hash aggregate in the stage, its CASE WHENs split at the
default threshold, and their methods take the buffer row. The test now asserts
both. The parameter-limit test now asserts that the same CASE WHEN splits under
the default limit, and the discarded-pass test asserts the split. Done in
1d4975551c2.
--
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]
