HyukjinKwon commented on code in PR #59069:
URL: https://github.com/apache/spark/pull/59069#discussion_r4118196621
##########
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:
Following up on the thread above: I think there is one aggregate landing
point where `INPUT_ROW` is *not* in scope, a split common subexpression.
In the keyed `HashAggregateExec` update, `ctx.INPUT_ROW = unsafeRowBuffer`
is set before
`subexpressionEliminationForWholeStageCodegen(boundUpdateExprs.flatten)`.
`unsafeRowBuffer` is a local of the consume method (`UnsafeRow
agg_unsafeRowAggBuffer = null;`). If the CASE WHEN is itself a common
subexpression and the subexpression code is over the split threshold, the split
pass regenerates it inside
```java
private void subExpr_N(<inputVars from getLocalInputVariableValues>) { ... }
```
Those parameters never include `INPUT_ROW`. With a `CommonExpressionRef` in
a branch, this arm forces `needsRow`, so the CASE WHEN's split call inside
`subExpr_N` becomes `caseWhen_0(..., agg_unsafeRowAggBuffer)`. That names an
undeclared variable, so the stage fails to compile: an internal error under
testing, and in production a fallback to non-codegen execution. Before this PR
the CASE WHEN wasn't split there, and a short definition like `nullif(v * 3,
1)` is inlined by `CommonExprSlots` without touching the row, so the same query
compiled.
Untested repro shape, combining two of your tests (the `nullif` branch from
"a common expression in a branch under a keyed aggregate takes the buffer",
used twice so it becomes a common subexpression):
```scala
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 cw = s"CASE $branches ELSE 0 END"
sql(s"SELECT k, sum($cw), max($cw) FROM t GROUP BY k")
```
The existing split-common-subexpression test runs under a projection, where
`INPUT_ROW` is null, so it doesn't reach this. If this reproduces, one option
is to refuse the split (return `None`) when `needsRow` is set while generating
inside a subexpression function, or have `subExpr_N` take `INPUT_ROW` when it
is set. Either way a test for this shape would pin it.
--
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]
