LuciferYang commented on code in PR #58881:
URL: https://github.com/apache/spark/pull/58881#discussion_r4466842384
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -1322,23 +1330,26 @@ class CodegenContext extends Logging {
val args = mutable.LinkedHashMap.empty[String, VariableValue]
val inputsToEvaluate = mutable.ArrayBuffer.empty[ExprCode]
// Whether the code can read `v` from the method: false for a value no
parameter can carry
- // where the policy asks for names. `ExpandExec` hands out a
`VariableValue` naming a slot of a
- // compacted mutable state array, and a `SimpleExprValue` is an expression
rather than a name
+ // where the policy asks for names. A `SimpleExprValue` is an expression
rather than a name
// -- `posexplode_outer` gives its position the nullness `index == -1`,
and a `Byte` or `Short`
// literal's value is `(byte)1`. A field or a literal needs no parameter
and is read as it
- // stands; so is a field handed out as a `VariableValue` where the policy
says so, passing one
- // being refused by the split methods' own check, a slot of a compacted
array included.
+ // stands; so is a `VariableValue` naming a field `addMutableState`
registered, where the
+ // policy says so, passing one being refused by the whole stage split's
own check (`splitCall`)
+ // in tests. A slot of a compacted array is not a name at all, so it has
to be handed out as a
+ // `GlobalValue`, and `INPUT_ROW` cannot be one. Testing checks for one
handed out as a
+ // variable: here, where the policy asks for names, for what the walk
reaches before its first
+ // refusal; elsewhere, in the method taking it (`assertMethodParameters`).
The row a whole stage
+ // split appends is not checked (SPARK-60510).
def take(v: ExprValue): Boolean = v match {
case local: VariableValue =>
val name = local.variableName
- val isName = name.nonEmpty &&
Character.isJavaIdentifierStart(name.head) &&
- name.forall(Character.isJavaIdentifierPart)
- // A slot of a compacted array is a field under every policy, since no
parameter can be
- // named `array[i]`, where the whole stage split is on; with it off,
every policy answers
- // as it did before that split existed. Another field is passed where
the policy says so.
- val isField = (isCompactedSlot(name) && wholeStageSplit !=
WholeStageSplit.Off) ||
- (policy.fieldsNeedNoArgument && mutableStateNames.contains(name))
- val taken = isName || isField || !policy.onlyNames
+ val isName = isJavaName(name)
+ if (Utils.isTesting && policy.onlyNames) {
+ assert(isName, s"$name cannot name a parameter of a method this code
moves into; " +
+ nameFix(name))
+ }
+ val isField = policy.fieldsNeedNoArgument &&
mutableStateNames.contains(name)
+ val taken = isName || !policy.onlyNames
Review Comment:
Done in 97c803e1beb: `CodegenContext.checkParameterNames` defaults to
`Utils.isTesting`, and `collectInputs` and `assertParametersAreNames` use it.
With it off, `CollectInputsSuite` checks that the whole-stage split and common
expression policies refuse the method for a slot and for an `INPUT_ROW` that is
a slot, and `WithExpressionEvalSuite` checks the inline body again with
`markersIn(...) == 4`. With `val taken = true`, those three tests fail.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/ExpandExec.scala:
##########
@@ -189,9 +189,15 @@ case class ExpandExec(
"resultValue",
v => s"$v = ${CodeGenerator.defaultValue(firstExpr.dataType)};")
+ // `addMutableState` hands back a field: a name of its own, or a
slot in a compacted array
+ // (`mutableStateArray_0[3]`) for a state not inlined. Either is
reachable from every
+ // method of the generated class and is not a local variable, so a
consumer collecting
+ // local variables to pass into a split function must not take it
for one -- a slot
+ // expression is not a legal parameter name. Whole-stage
subexpression elimination wraps
+ // its own mutable state the same way.
Review Comment:
Done in 97c803e1beb: the `@return` of `addMutableState` now says to hand the
result out with `JavaCode.global` or `isNullGlobal`.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -1411,13 +1422,15 @@ class CodegenContext extends Logging {
}
}
- /**
- * Whether `name` is a slot of a compacted mutable state array, such as
`mutableStateArray_0[3]`,
- * which `addMutableState` hands out for a state it does not inline: a field
of the class.
- */
- private def isCompactedSlot(name: String): Boolean = name match {
- case CodeGenerator.CompactedSlot(array) =>
mutableStateNames.contains(array)
- case _ => false
+ /** Whether `name` can name a parameter: a slot of a compacted array,
`array[i]`, cannot. */
+ private def isJavaName(name: String): Boolean =
+ SourceVersion.isIdentifier(name) && !SourceVersion.isKeyword(name)
+
+ /** How the producer of `name`, a value that cannot name a parameter, can
hand it out instead. */
+ private def nameFix(name: String): String = if (name == INPUT_ROW) {
+ "set INPUT_ROW to a local or a field added with forceInline = true"
+ } else {
+ "hand out a field with JavaCode.global or isNullGlobal"
}
Review Comment:
Applied in 97c803e1beb, with an assertion on the hint for `false` in
`CollectInputsSuite`.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/WholeStageCodegenSuite.scala:
##########
@@ -1932,6 +1962,55 @@ class WholeStageCodegenSuite extends SharedSparkSession
query.collect().toSeq
}
+ /**
+ * Checks the query `query` builds from a FROM clause that UNPIVOTs two
string columns into
+ * `v`, with the whole stage split on and off: a method named like `method`
calls `lpad` on the
+ * slot of the compacted mutable state array that `ExpandExec` holds `v` in,
`check` holds of
+ * the code, the code compiles, and the rows match those computed without
whole stage codegen.
+ * With the slot handed out as a local variable instead, building that
method fails the testing
+ * check on its parameters.
+ */
+ private def checkExpandSlotRead(
+ method: String,
+ query: String => String,
+ check: Seq[CodeAndComment] => Unit = _ => ()): Unit = {
+ withTempView("t") {
+ spark.range(10).selectExpr("concat('a', CAST(id AS STRING)) AS a",
+ "concat('b', CAST(id AS STRING)) AS b").createOrReplaceTempView("t")
+ def df: DataFrame = sql(query("FROM t UNPIVOT (v FOR c IN (a, b))"))
+ val expected = withoutWholeStage(df)
+ Seq("true", "false").foreach { split =>
+ withSQLConf(SQLConf.WHOLESTAGE_SPLIT_EXPRESSIONS.key -> split) {
+ val code = genCode(df)
+ // A subexpression method or a switch case also writes into a slot
of the same array,
+ // so the array's name alone does not show a read.
+ assert(code.exists(c => methodsHolding(c,
"].lpad(").exists(_.contains(method))),
+ s"$method $split")
+ withClue(s"$method $split: ")(check(code))
+ code.foreach(CodeGenerator.compile)
+ checkAnswer(df, expected)
+ }
+ }
+ withSQLConf(SQLConf.WHOLESTAGE_SPLIT_EXPRESSIONS.key -> "false") {
+ val e = intercept[AssertionError](genCodeWithLocalExpandOutput(df))
+ assert(s"\\w*$method\\w* takes
\\w+\\[\\d+\\]".r.findFirstIn(e.getMessage).isDefined,
+ e.getMessage)
+ }
+ }
+ }
+
+ /**
+ * `genCode`, with the bottom `ExpandExec` of the plan handing out its
output as local variables,
+ * as it did before SPARK-59589.
+ */
+ private def genCodeWithLocalExpandOutput(df: DataFrame): Seq[CodeAndComment]
=
Review Comment:
Done in 97c803e1beb: `genCode` takes the change of the plan, and the helper
is gone.
--
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]
