[SPARK-22226][SQL] splitExpression can create too many method calls in the outer class

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

@@ -77,6 +77,20 @@ case class SubExprEliminationState(isNull: String, value: String)
  */
 case class SubExprCodes(codes: Seq[String], states: Map[Expression, SubExprEliminationState])
 
+/**
+ * The main information about a new added function.
+ *
+ * @param functionName String representing the name of the function
+ * @param subclassName Optional value which is empty if the function is added to
+ *                     the outer class, otherwise it contains the name of the
+ *                     inner class in which the function has been added.
+ * @param subclassInstance Optional value which is empty if the function is added to
+ *                         the superclass, otherwise it contains the name of the
+ *                         instance of the inner class in the outer class.
+ */
+private[codegen] case class NewFunction(functionName: String, subclassName: Option[String],
+  subclassInstance: Option[String])
+
 /**
  * A context for codegen, tracking a list of objects that could be passed into generated Java
  * function.
@@ -277,13 +291,25 @@ class CodegenContext {
       funcName: String,
       funcCode: String,
       inlineToOuterClass: Boolean = false): String = {
+    val newFunction = addNewFunctionInternal(funcName, funcCode, inlineToOuterClass)
+    newFunction match {
+      case NewFunction(functionName, None, None) => functionName
+      case NewFunction(functionName, Some(_), Some(subclassInstance)) =>
+        subclassInstance + "." + functionName
+    }
+  }
+
+  private[this] def addNewFunctionInternal(
+      funcName: String,
+      funcCode: String,
+      inlineToOuterClass: Boolean): NewFunction = {
     // The number of named constants that can exist in the class is limited by the Constant Pool
     // limit, 65,536. We cannot know how many constants will be inserted for a class, so we use a
-    // threshold of 1600k bytes to determine when a function should be inlined to a private, nested
+    // threshold of 1000k bytes to determine when a function should be inlined to a private, nested
     // sub-class.
     val (className, classInstance) = if (inlineToOuterClass) {
       outerClassName -> ""
-    } else if (currClassSize > 1600000) {
+    } else if (currClassSize > 1000000) {
       val className = freshName("NestedClass")
       val classInstance = freshName("nestedClassInstance")
 
@@ -294,17 +320,23 @@ class CodegenContext {
       currClass()
     }
 
-    classSize(className) += funcCode.length
-    classFunctions(className) += funcName -> funcCode
+    addNewFunctionToClass(funcName, funcCode, className)
 
     if (className == outerClassName) {
-      funcName
+      NewFunction(funcName, None, None)
     } else {
-
-      s"$classInstance.$funcName"
+      NewFunction(funcName, Some(className), Some(classInstance))
     }
   }
 
+  private[this] def addNewFunctionToClass(
+      funcName: String,
+      funcCode: String,
+      className: String) = {
+    classSize(className) += funcCode.length
+    classFunctions(className) += funcName -> funcCode
+  }
+
   /**
    * Declares all function code. If the added functions are too many, split them into nested
    * sub-classes to avoid hitting Java compiler constant pool limitation.
@@ -798,10 +830,35 @@ class CodegenContext {
            |  ${makeSplitFunction(body)}
            |}
          """.stripMargin
-        addNewFunction(name, code)
+        addNewFunctionInternal(name, code, inlineToOuterClass = false)
       }
 
-      foldFunctions(functions.map(name => s"$name(${arguments.map(_._2).mkString(", ")})"))
+      val outerClassFunctions = functions
+        .filter(_.subclassName.isEmpty)
+        .map(_.functionName)
+
+      val innerClassFunctions = functions
+        .filter(_.subclassName.isDefined)
+        .foldLeft(Map.empty[(String, String), Seq[String]]) { case (acc, f) =>
+          val key = (f.subclassName.get, f.subclassInstance.get)
+          acc.updated(key, acc.getOrElse(key, Seq.empty[String]) ++ Seq(f.functionName))
+        }
+        .map { case ((subclassName, subclassInstance), subclassFunctions) =>
+          // Adding a new function to each subclass which contains
+          // the invocation of all the ones which have been added to
+          // that subclass
+          val code = s"""
+              |private $returnType $func($argString) {
+              |  ${makeSplitFunction(foldFunctions(subclassFunctions.map(name =>
+                    s"$name(${arguments.map(_._2).mkString(", ")})")))}
+              |}
+            """.stripMargin
+          addNewFunctionToClass(func, code, subclassName)
+          s"$subclassInstance.$func"
+        }
+
+      foldFunctions((outerClassFunctions ++ innerClassFunctions).map(
+        name => s"$name(${arguments.map(_._2).mkString(", ")})"))
     }
   }
 

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

@@ -201,6 +201,23 @@ class CodeGenerationSuite extends SparkFunSuite with ExpressionEvalHelper {
     }
   }
 
+  test("SPARK-22226: group splitted expressions into one method per nested class") {
+    val length = 10000
+    val expressions = Seq.fill(length) {
+      ToUTCTimestamp(
+        Literal.create(Timestamp.valueOf("2017-10-10 00:00:00"), TimestampType),
+        Literal.create("PST", StringType))
+    }
+    val plan = GenerateMutableProjection.generate(expressions)
+    val actual = plan(new GenericInternalRow(length)).toSeq(expressions.map(_.dataType))
+    val expected = Seq.fill(length)(
+      DateTimeUtils.fromJavaTimestamp(Timestamp.valueOf("2017-10-10 07:00:00")))
+
+    if (actual != expected) {
+      fail(s"Incorrect Evaluation: expressions: $expressions, actual: $actual, expected: $expected")
+    }
+  }
+
   test("test generated safe and unsafe projection") {
     val schema = new StructType(Array(
       StructField("a", StringType, true),