HyukjinKwon commented on code in PR #59225:
URL: https://github.com/apache/spark/pull/59225#discussion_r4175100231
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -1813,6 +2364,89 @@ object CodeGenerator extends Logging {
// beyond which methods will be rejected from JIT compilation
final val DEFAULT_JVM_HUGE_METHOD_LIMIT = 8000
+ /** A slot of a compacted mutable state array, such as
`mutableStateArray_0[3]`. */
+ private[codegen] val CompactedSlot = "([A-Za-z_$][\\w$]*)\\[\\d+\\]".r
+
+ /** A name and an opening parenthesis: a call, or a keyword such as `if`. */
+ private[codegen] val CallSite = "(?<![\\w$])([A-Za-z_$][\\w$]*)\\s*\\(".r
+
+ /**
+ * The comment marking the code of the expression of index `i` in
+ * `CodegenContext.wholeStageSplitsRecorded` under `WholeStageSplit.Record`.
+ */
+ private[sql] def wholeStageSplitMarker(i: Int): String =
s"$WholeStageSplitMarkerStart$i*/"
+
+ private val WholeStageSplitMarkerStart = "/*wholeStageSplit:"
+
+ /** Whether `code` holds a `wholeStageSplitMarker`. */
+ private[codegen] def holdsWholeStageSplitMarker(code: String): Boolean =
+ code.contains(WholeStageSplitMarkerStart)
+
+ private val WholeStageSplitMarker = """/\*wholeStageSplit:(\d+)\*/""".r
+
+ /** `source` without the comments `wholeStageSplitMarker` adds. */
+ private[sql] def stripWholeStageSplitMarkers(source: String): String =
+ WholeStageSplitMarker.replaceAllIn(source, "")
+
+ /**
+ * The methods of the generated Java `source` that hold each
`wholeStageSplitMarker`, by its
+ * index. A method is a body opened, in the body of a class, after a
parameter list and any
+ * `throws` clause; a block, a lambda or an anonymous class inside a method
belongs to that
+ * method. String and character literals and comments are skipped, apart
from the markers.
+ */
+ private[sql] def wholeStageSplitMarkersByMethod(source: String): Map[Int,
Set[String]] = {
+ val found = mutable.HashMap.empty[Int, Set[String]]
+ // Each open brace's scope: Some(name) for a method, None for a class body
or a block, with
+ // whether it is a class body.
+ val scopes = mutable.Stack[(Boolean, Option[String])]((true, None))
+ var start = 0
+ var i = 0
+ def method: Option[String] = scopes.collectFirst { case (_, name @
Some(_)) => name }.flatten
+ while (i < source.length) {
+ source.charAt(i) match {
+ case '"' | '\'' =>
+ val quote = source.charAt(i)
+ i += 1
+ while (i < source.length && source.charAt(i) != quote) {
+ if (source.charAt(i) == '\\') i += 1
+ i += 1
+ }
+ case '/' if source.startsWith("//", i) =>
+ i = source.indexOf('\n', i) match { case -1 => source.length; case e
=> e }
+ case '/' if source.startsWith("/*", i) =>
+ val end = source.indexOf("*/", i + 2) match { case -1 =>
source.length; case e => e }
+ WholeStageSplitMarker.findPrefixMatchOf(source.substring(i,
math.min(end + 2,
+ source.length))).foreach { m =>
+ val index = m.group(1).toInt
+ found(index) = found.getOrElse(index, Set.empty) ++ method
+ }
+ i = end + 1
+ case '{' =>
+ val header = source.substring(start, i).trim
+ val isClass =
"\\b(class|interface|enum)\\b".r.findFirstIn(header).isDefined
Review Comment:
**Nit (P3):** `"\\b(class|interface|enum)\\b".r` is compiled on every `{`
while scanning the (potentially megabyte-sized) generated source, whereas
`CompactedSlot`, `CallSite`, `MethodHeader`, and `WholeStageSplitMarker` are
all hoisted to module vals. Consider hoisting this one too so it compiles once.
**Verification:**
- **Regression:** Existing wholeStageSplitMarkersByMethod tests still pass.
- **Inspection:** No per-iteration .r compilation remains.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -1271,49 +1696,158 @@ class CodegenContext extends Logging {
DEFAULT_JVM_HUGE_METHOD_LIMIT / 2 / (MAX_BYTES_PER_SPLIT_CALL + 2 *
arguments.length))
/**
- * The calls to the functions `splitExpressions` split out, grouped into
functions of their own
- * where there are too many for the calling method to be JIT-compiled, level
by level.
+ * The calls to the functions `splitExpressions` split out, grouped by
`groupSplitCalls` where
+ * there are too many for the calling method to be JIT-compiled. Every call
takes `arguments`, so
+ * the groups are of `SPLIT_CALLS_PER_GROUP` calls, and the calls stay in
the calling method while
+ * they fit in `splitCallsPerMethod`.
+ */
+ private def groupFunctionCalls(
+ calls: Seq[String],
+ funcName: String,
+ arguments: Seq[(String, String)],
+ returnType: String,
+ makeSplitFunction: String => String,
+ foldFunctions: Seq[String] => String): Seq[String] = {
+ groupSplitCalls(calls.map(SplitCall(_, arguments)), funcName, returnType,
makeSplitFunction,
+ foldFunctions, callsLeft = splitCallsPerMethod(arguments)).map(_.call)
+ }
+
+ /** A call to a split function, and the (type, name) arguments it passes. */
+ 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.
+ * Where `audited`, its arguments are kept under testing for
`assertSplitCallsWithin`: those of a
+ * whole stage split, computed from what its code reads where it was
generated, which an operator
+ * moving that code into a method of its own must take. The other splits'
arguments are their
+ * callers' own, which may be locals of the method they are made in.
+ */
+ private def splitCall(
+ name: String,
+ arguments: Seq[(String, String)],
+ returnType: String,
+ body: String,
+ audited: Boolean = false): 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 argDefinitionString = arguments.map { case (t, argument) => s"$t
$argument" }.mkString(", ")
+ val code =
+ s"""
+ |private $returnType $name($argDefinitionString) {
+ | $body
+ |}
+ """.stripMargin
+ val call = s"${addNewFunction(name,
code)}(${arguments.map(_._2).mkString(", ")})"
+ if (INPUT_ROW != null && arguments.exists(_._2 == INPUT_ROW)) {
+ functionsTakingRow += name
+ }
+ if (audited && Utils.isTesting) {
+ splitCallArguments(name) = arguments.map(_._2)
+ }
+ SplitCall(call, arguments)
+ }
+
+ /**
+ * The calls to split functions, grouped into functions of their own, level
by level, until at
+ * most `callsLeft` remain in the calling method.
*
* The calls are what is left in the calling method, one per block of code,
and a `CASE WHEN`
* of a thousand branches has hundreds of blocks: their calls alone take the
caller past the
* 8000 bytes HotSpot compiles
(`CodeGenerator.DEFAULT_JVM_HUGE_METHOD_LIMIT`), so the method
* every row goes through runs interpreted. A call and its fold compile to
at most about 20 bytes
* plus 2 for each argument - the invocation, a load per argument, and the
fold's store, test
- * and branch - so the calls are left as they are while they fit in half of
that limit, the rest
- * being the caller's own code, and grouped otherwise. The groups are small,
of
- * `CodeGenerator.SPLIT_CALLS_PER_GROUP` calls, so that C2 can inline the
split functions into
- * each group within its inlining budget; a group of hundreds of calls stays
compilable but has
- * most of its calls left out of inlining, and runs slower. A group is
folded with
- * `foldFunctions` and wrapped with `makeSplitFunction`, which is what
- * `generateInnerClassesFunctionCalls` does with the functions of one inner
class, so every
- * caller of `splitExpressions` supports it.
- */
- private def groupFunctionCalls(
- calls: Seq[String],
+ * and branch - so outside whole stage codegen the calls are left as they
are while they fit in
+ * half of that limit, the rest being the caller's own code
(`splitCallsPerMethod`). Under whole
+ * stage codegen the caller is the operator's method, which holds every
expression's calls, so
+ * at most `SPLIT_CALLS_PER_GROUP` are left there for each.
+ *
+ * The groups are small, of `CodeGenerator.SPLIT_CALLS_PER_GROUP` calls, so
that C2 can inline
+ * the split functions into each group within its inlining budget; a group
of hundreds of calls
+ * stays compilable but has most of its calls left out of inlining, and runs
slower. A group
+ * takes the union of what its calls take, and is closed when one more call
would take it past
+ * those calls or past the parameters a method can have; where no group can
hold two calls, the
+ * calls are left as they are. A group is folded with `foldFunctions` and
wrapped with
+ * `makeSplitFunction`, which is what `generateInnerClassesFunctionCalls`
does with the functions
+ * of one inner class, so every caller of `splitExpressions` supports it.
+ */
+ private def groupSplitCalls(
+ calls: Seq[SplitCall],
funcName: String,
- arguments: Seq[(String, String)],
returnType: String,
makeSplitFunction: String => String,
foldFunctions: Seq[String] => String,
- level: Int = 0): Seq[String] = {
- val callsPerMethod = splitCallsPerMethod(arguments)
- if (calls.length <= callsPerMethod) {
+ callsLeft: Int,
+ level: Int = 0,
+ audited: Boolean = false): Seq[SplitCall] = {
+ if (calls.length <= callsLeft) {
calls
} else {
- val argDefinitionString = arguments.map { case (t, name) => s"$t $name"
}.mkString(", ")
- val argInvocationString = arguments.map(_._2).mkString(", ")
- val groupCalls = calls.grouped(SPLIT_CALLS_PER_GROUP).zipWithIndex.map {
case (group, i) =>
- val name = s"${funcName}_group${level}_$i"
- val code =
- s"""
- |private $returnType $name($argDefinitionString) {
- | ${makeSplitFunction(foldFunctions(group))}
- |}
- """.stripMargin
- s"${addNewFunction(name, code)}($argInvocationString)"
- }.toSeq
- groupFunctionCalls(groupCalls, funcName, arguments, returnType,
makeSplitFunction,
- foldFunctions, level + 1)
+ 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)), audited)
+ }.toSeq
+ groupSplitCalls(grouped, funcName, returnType, makeSplitFunction,
foldFunctions,
+ callsLeft, level + 1, audited)
+ }
+ }
+ }
+
+ /** 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
+
+ /**
+ * Under testing, the names every split function's calls pass, by the
function's name, for
+ * `assertSplitCallsWithin`. A function `removeFunctionsSince` removes
leaves it too.
Review Comment:
**Nit (P3):** The clause "removes leaves it too" doesn't parse. It looks
meant to say that `removeFunctionsSince` also removes a removed function's
entry from `splitCallArguments` (which its body does); consider rewording.
**Verification:**
- **Inspection:** The reworded comment is grammatical and accurate.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -3069,6 +3069,37 @@ object SQLConf {
.booleanConf
.createWithDefault(true)
+ val WHOLESTAGE_SPLIT_EXPRESSIONS =
+ buildConf("spark.sql.codegen.wholeStage.splitExpressions")
+ .internal()
+ .doc("When true, whole stage codegen splits the generated code of an
expression that " +
+ "supports it, such as a CASE WHEN with many branches, into methods
that take the input " +
+ "variables they read as parameters, the way code generation outside
whole stage codegen " +
+ "splits it; and when subexpression elimination discards the code of
its first pass for " +
+ "its split one, it removes the methods that pass added. When false,
the code stays in " +
+ "the method of its operator, where a large enough expression goes past
the JVM's 64KB " +
+ "method limit and fails to compile, and the class keeps those
methods.")
+ .version("4.4.0")
+ .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
+ .booleanConf
+ .createWithDefault(true)
+
+ val WHOLESTAGE_SPLIT_EXPRESSIONS_METHOD_LIMIT =
+ buildConf("spark.sql.codegen.wholeStage.splitExpressions.methodLimit")
+ .internal()
+ .doc("The largest method, in bytes of bytecode, a whole stage keeps
unsplit when " +
+ "spark.sql.codegen.wholeStage.splitExpressions is true. The stage's
code is first " +
+ "generated with no expression split and compiled; only when that fails
or a method " +
+ "is past this size is it generated again with the expressions split,
and the split " +
+ "code is kept when it makes the stage's largest method smaller. The
default is " +
Review Comment:
**Nit (P3):** The doc says the split code is kept "when it makes the stage's
largest method smaller", but `chooseCode` keeps it when `bytesOver(split) <
over.values.sum` — i.e. when the *total* bytecode of the methods past the limit
goes down, not the single largest method (the two can disagree when several
methods exceed the limit). `chooseCode`'s own scaladoc states this correctly;
consider aligning the config doc.
**Verification:**
- **Inspection:** The doc matches chooseCode's bytesOver(split) <
over.values.sum condition.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -1915,6 +2549,26 @@ object CodeGenerator extends Logging {
throw e.getCause
}
+ /**
+ * The class loader the compile cache keys a class on, held weakly: the
calling thread's context
+ * class loader, or Spark's.
+ */
+ private[sql] def classLoaderKey(): AnyRef =
+ new HashableWeakReference(Utils.getContextOrSparkClassLoader)
+
+ /** Drops `code`'s class from the compile cache: a variant whole stage
codegen did not keep. */
+ private[sql] def invalidateCompiled(code: CodeAndComment): Unit = {
+ cache.invalidate(
+ (new HashableWeakReference(Utils.getContextOrSparkClassLoader),
CodeCompiler.active(code),
+ code))
+ }
+
+ /** Whether `code`'s class is in the compile cache, for tests. */
+ private[sql] def isCompiled(code: CodeAndComment): Boolean =
Review Comment:
**Nit (P3):** `isCompiled` has no caller anywhere in the tree (only
`invalidateCompiled` is used, by `trialCompile`). The "for tests" comment
suggests a test was intended; either wire it up or drop the helper.
**Verification:**
- **Compile:** The module compiles after the change.
- **Inspection:** No remaining reference to isCompiled, or a new test uses
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]
