Refactor ctx.splitExpressions().

Author: ueshin

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala

@@ -610,8 +610,13 @@ class CodegenContext {
     splitExpressions(expressions, "apply", ("InternalRow", row) :: Nil)
   }
 
-  private def splitExpressions(
-      expressions: Seq[String], funcName: String, arguments: Seq[(String, String)]): String = {
+  def splitExpressions(
+      expressions: Seq[String],
+      funcName: String,
+      arguments: Seq[(String, String)],
+      returnType: String = "void",
+      makeSplitFunction: String => String = identity,
+      foldFunctions: Seq[String] => String = _.mkString("", ";\n", ";")): String = {
     val blocks = new ArrayBuffer[String]()
     val blockBuilder = new StringBuilder()
     for (code <- expressions) {
@@ -632,18 +637,19 @@ class CodegenContext {
       blocks.head
     } else {
       val func = freshName(funcName)
+      val argString = arguments.map { case (t, name) => s"$t $name" }.mkString(", ")
       val functions = blocks.zipWithIndex.map { case (body, i) =>
         val name = s"${func}_$i"
         val code = s"""
-           |private void $name(${arguments.map { case (t, name) => s"$t $name" }.mkString(", ")}) {
-           |  $body
+           |private $returnType $name($argString) {
+           |  ${makeSplitFunction(body)}
            |}
          """.stripMargin
         addNewFunction(name, code)
         name
       }
 
-      functions.map(name => s"$name(${arguments.map(_._2).mkString(", ")});").mkString("\n")
+      foldFunctions(functions.map(name => s"$name(${arguments.map(_._2).mkString(", ")})"))
     }
   }
 

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateOrdering.scala

@@ -72,7 +72,7 @@ object GenerateOrdering extends CodeGenerator[Seq[SortOrder], Ordering[InternalR
    * Generates the code for ordering based on the given order.
    */
   def genComparisons(ctx: CodegenContext, ordering: Seq[SortOrder]): String = {
-    def comparisons(orderingGroup: Seq[SortOrder]) = orderingGroup.map { order =>
+    val comparisons = ordering.map { order =>
       val eval = order.child.genCode(ctx)
       val asc = order.isAscending
       val isNullA = ctx.freshName("isNullA")
@@ -117,43 +117,31 @@ object GenerateOrdering extends CodeGenerator[Seq[SortOrder], Ordering[InternalR
             }
           }
       """
-    }.mkString("\n")
-
-    /*
-     * 40 = 7000 bytes / 170 (around 170 bytes per ordering comparison).
-     * The maximum byte code size to be compiled for HotSpot is 8000 bytes.
-     * We should keep less than 8000 bytes.
-     */
-    val numberOfComparisonsThreshold = 40
-
-    if (ordering.size <= numberOfComparisonsThreshold) {
-      comparisons(ordering)
-    } else {
-      val groupedOrderingItr = ordering.grouped(numberOfComparisonsThreshold)
-      val funcNamePrefix = ctx.freshName("compare")
-      val funcNames = groupedOrderingItr.zipWithIndex.map { case (orderingGroup, i) =>
-        val funcName = s"${funcNamePrefix}_$i"
-        val funcCode =
-          s"""
-             |private int $funcName(InternalRow a, InternalRow b) {
-             |  InternalRow ${ctx.INPUT_ROW} = null;  // Holds current row being evaluated.
-             |  ${comparisons(orderingGroup)}
-             |  return 0;
-             |}
-          """.stripMargin
-        ctx.addNewFunction(funcName, funcCode)
-        funcName
-      }
+    }
 
-      funcNames.zipWithIndex.map { case (funcName, i) =>
+    ctx.splitExpressions(
+      expressions = comparisons,
+      funcName = "compare",
+      arguments = Seq(("InternalRow", "a"), ("InternalRow", "b")),
+      returnType = "int",
+      makeSplitFunction = { body =>
         s"""
-           |int comp_$i = ${funcName}(a, b);
-           |if (comp_$i != 0) {
-           |  return comp_$i;
-           |}
-        """.stripMargin
-      }.mkString
-    }
+          InternalRow ${ctx.INPUT_ROW} = null;  // Holds current row being evaluated.
+          $body
+          return 0;
+        """
+      },
+      foldFunctions = { funCalls =>
+        val comp = ctx.freshName("comp")
+        funCalls.zipWithIndex.map { case (funCall, i) =>
+          s"""
+            int ${comp}_$i = $funCall;
+            if (${comp}_$i != 0) {
+              return ${comp}_$i;
+            }
+          """
+        }.mkString
+      })
   }
 
   protected def create(ordering: Seq[SortOrder]): BaseOrdering = {

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/OrderingSuite.scala

@@ -137,7 +137,7 @@ class OrderingSuite extends SparkFunSuite with ExpressionEvalHelper {
     // this is FAILING prior to SPARK-16845, but it should be passing after SPARK-16845
     GenerateOrdering.generate(Array.fill(450)(sortOrder))
 
-    // verify that we can support up to 10000 ordering comparisons, which should be sufficient
-    GenerateOrdering.generate(Array.fill(10000)(sortOrder))
+    // verify that we can support up to 5000 ordering comparisons, which should be sufficient
+    GenerateOrdering.generate(Array.fill(5000)(sortOrder))
   }
 }