dongjoon-hyun commented on code in PR #59042:
URL: https://github.com/apache/spark/pull/59042#discussion_r4108398490
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -1248,7 +1248,67 @@ class CodegenContext extends Logging {
makeSplitFunction,
foldFunctions)
- foldFunctions(outerClassFunctionCalls ++ innerClassFunctionCalls)
+ foldFunctions(groupFunctionCalls(
+ (outerClassFunctionCalls ++ innerClassFunctionCalls).toSeq,
+ func,
+ arguments,
+ returnType,
+ makeSplitFunction,
+ foldFunctions))
+ }
+ }
+
+ /**
+ * 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 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.
+ */
+ /**
+ * How many calls to split functions taking `arguments` fit in one method
with room to spare
+ * below the 8000 bytes HotSpot compiles; see `groupFunctionCalls`.
+ */
+ private def splitCallsPerMethod(arguments: Seq[(String, String)]): Int =
math.max(2,
Review Comment:
This budget assumes a single `splitExpressions` call owns half of the
8000-byte limit. When one method invokes `splitExpressions` multiple times
(e.g. `ConcatWs` calls it three times in the same method), the combined calls
can still push the caller past the limit. It's probably rare, but could we
document this limitation here?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -1248,7 +1248,67 @@ class CodegenContext extends Logging {
makeSplitFunction,
foldFunctions)
- foldFunctions(outerClassFunctionCalls ++ innerClassFunctionCalls)
+ foldFunctions(groupFunctionCalls(
+ (outerClassFunctionCalls ++ innerClassFunctionCalls).toSeq,
+ func,
+ arguments,
+ returnType,
+ makeSplitFunction,
+ foldFunctions))
+ }
+ }
+
+ /**
+ * The calls to the functions `splitExpressions` split out, grouped into
functions of their own
Review Comment:
This Scaladoc is immediately followed by the Scaladoc of
`splitCallsPerMethod`, so it is orphaned and not attached to any member. Could
you move it right above `private def groupFunctionCalls(`?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -1314,13 +1374,28 @@ class CodegenContext extends Logging {
val argDefinitionString = arguments.map { case (t, name) => s"$t $name"
}.mkString(", ")
val argInvocationString = arguments.map(_._2).mkString(", ")
+ val callsPerMethod = splitCallsPerMethod(arguments)
innerClassToFunctions.flatMap {
case ((innerClassName, innerClassInstance), innerClassFunctions) =>
// for performance reasons, the functions are prepended, instead of
appended,
// thus here they are in reversed order
val orderedFunctions = innerClassFunctions.reverse
- if (orderedFunctions.size > MERGE_SPLIT_METHODS_THRESHOLD) {
+ if (orderedFunctions.size > callsPerMethod) {
+ // Too many calls for one method to stay JIT-compilable: one merged
method per part,
+ // whose calls the caller groups further if there are still too many
of them.
+ orderedFunctions.grouped(callsPerMethod).zipWithIndex.map { case
(part, i) =>
Review Comment:
Each `_part` method here holds `callsPerMethod` calls (~166-200 with 1-2
arguments). However, the PR description says groups of ~180 calls make C2
refuse most inlining and run ~35% slower, which is why `SPLIT_CALLS_PER_GROUP`
is 8. So this path re-creates the size the PR is trying to avoid. Since merging
inside the inner class is for the outer class' constant pool, how about
splitting into parts of `SPLIT_CALLS_PER_GROUP` here and letting
`groupFunctionCalls` handle further levels? If the current size is intentional,
could you add a comment explaining why?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -1248,7 +1248,67 @@ class CodegenContext extends Logging {
makeSplitFunction,
foldFunctions)
- foldFunctions(outerClassFunctionCalls ++ innerClassFunctionCalls)
+ foldFunctions(groupFunctionCalls(
+ (outerClassFunctionCalls ++ innerClassFunctionCalls).toSeq,
+ func,
+ arguments,
+ returnType,
+ makeSplitFunction,
+ foldFunctions))
+ }
+ }
+
+ /**
+ * 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 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.
+ */
+ /**
+ * How many calls to split functions taking `arguments` fit in one method
with room to spare
+ * below the 8000 bytes HotSpot compiles; see `groupFunctionCalls`.
+ */
+ private def splitCallsPerMethod(arguments: Seq[(String, String)]): Int =
math.max(2,
+ DEFAULT_JVM_HUGE_METHOD_LIMIT / 2 / (MAX_BYTES_PER_SPLIT_CALL + 2 *
arguments.length))
+
+ private def groupFunctionCalls(
+ calls: Seq[String],
+ 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) {
+ 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)"
Review Comment:
When the outer class is past `GENERATED_CLASS_SIZE_THRESHOLD`,
`addNewFunction` can place a group function into an inner class. Then a group
in inner class A calls private methods of inner class B
(`innerInstance.f(...)`) or private split functions of the outer class. This is
legal Java, but I'm not sure this sibling-nested-class private access pattern
was exercised with Janino before. Could you confirm that one of the new tests
(e.g. the 30000-wide hash test) actually goes through this path?
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CodeGenerationSuite.scala:
##########
@@ -102,6 +102,91 @@ class CodeGenerationSuite extends SparkFunSuite with
ExpressionEvalHelper {
assert(actual.head == cases)
}
+ /**
+ * Runs `body` and fails if the code it generates has a method past the 8000
bytes HotSpot
+ * compiles, which `CodeGenerator` logs as it compiles the class.
+ */
+ private def assertJitCompilable(body: => Unit): Unit = {
+ val appender = new LogAppender("methods too long to be JIT compiled")
+ withLogAppender(appender,
+ loggerNames = Seq(classOf[CodeGenerator[_, _]].getName), level =
Some(Level.INFO)) {
+ body
+ }
+ val tooLong = appender.loggingEvents.map(_.getMessage.getFormattedMessage)
+ .filter(_.contains("too long to be JIT compiled"))
+ assert(tooLong.isEmpty, tooLong.mkString("\n"))
+ }
+
+ // The tests below pin SPARK-59783: an expression split into hundreds of
functions left one call
+ // per function in the method holding them, and the calls alone took that
method past the
+ // 8000 bytes HotSpot compiles, so every row ran it interpreted. Each test
covers one caller of
+ // `splitExpressions`, with as many children as it takes for that method to
cross the limit
+ // without the grouping: fewer where one child generates more code.
+
+ test("SPARK-59783: the calls to the functions CASE WHEN splits into are
JIT-compilable") {
+ val wide = 3000
+ val input = BoundReference(0, IntegerType, nullable = false)
+ val expression = CaseWhen(
+ (1 to wide).map(k => (EqualTo(input, Literal(k)), Multiply(input,
Literal(k)))), Literal(0))
+ assertJitCompilable {
+ val projection = GenerateMutableProjection.generate(Seq(expression))
+ Seq(7 -> 49, wide -> wide * wide, wide + 1 -> 0).foreach { case (in,
out) =>
+ assert(projection(new GenericInternalRow(Array[Any](in))).getInt(0) ==
out)
+ }
+ }
+ }
+
+ test("SPARK-59783: the calls to the functions COALESCE splits into are
JIT-compilable") {
+ val wide = 10000
+ val expression = Coalesce((0 until wide).map(i => BoundReference(i,
IntegerType, true)))
+ assertJitCompilable {
+ val projection = GenerateMutableProjection.generate(Seq(expression))
+ val values = new Array[Any](wide)
+ assert(projection(new GenericInternalRow(values)).isNullAt(0))
+ values(wide - 1) = 42
+ assert(projection(new GenericInternalRow(values)).getInt(0) == 42)
+ }
+ }
+
+ test("SPARK-59783: the calls to the functions IN splits into are
JIT-compilable") {
+ val wide = 10000
+ val input = BoundReference(0, IntegerType, nullable = false)
+ val expression = In(input, (1 to wide).map(k => Literal(k * 2)))
+ assertJitCompilable {
+ val projection = GenerateMutableProjection.generate(Seq(expression))
+ Seq(2 -> true, wide * 2 -> true, 3 -> false).foreach { case (in, out) =>
+ assert(projection(new
GenericInternalRow(Array[Any](in))).getBoolean(0) == out)
+ }
+ }
+ }
+
+ test("SPARK-59783: the calls to the functions a hash splits into are
JIT-compilable") {
Review Comment:
These tests compile very wide expressions with Janino (30000-wide hash
twice, 3000-wide ordering and `CASE WHEN`). Could you share how long the new
tests take? If they are slow, reducing the width to just past the threshold
might help.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -1740,6 +1815,19 @@ object CodeGenerator extends Logging {
// method which is going to be called by the outer class instead of the many
small ones
final val MERGE_SPLIT_METHODS_THRESHOLD = 3
+ /**
+ * The most bytecode a call to a split function and its fold take in the
calling method,
+ * besides a load per argument; see `CodegenContext.groupFunctionCalls`.
+ */
+ final val MAX_BYTES_PER_SPLIT_CALL = 20
Review Comment:
Is the upper bound of 20 bytes verified with `javap`? A call to an
inner-class function additionally needs `aload_0` + `getfield` for the
instance, and a fold like `CaseWhen`'s (store, compare, branch) gets close to
20 bytes. The half-limit margin probably covers it, but it would be good to
confirm.
--
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]
