Inline method calls when possible.

With this change, method calls are inlined only if the method being
called
  1. consists of operations supported by the UDF compiler, and
  2. is one of the folllowing:
    * a final method, or
    * a method in a final class, or
    * a method in a final class

Signed-off-by: Sean Lee <selee@nvidia.com>

docs/compatibility.md

@@ -384,3 +384,4 @@ When translating UDFs to Catalyst expressions, the supported UDF functions are l
 |                          | Array.empty[Float]                                       |
 |                          | Array.empty[Double]                                      |
 |                          | Array.empty[String]                                      |
+| Method call              | Only if the method being called <ol><li>consists of operations supported by the UDF compiler, and</li><li>is one of the folllowing:<ul><li>a final method, or</li><li>a method in a final class, or</li><li>a method in a final object</li></ul></li></ol> |

udf-compiler/src/main/scala/com/nvidia/spark/udf/CatalystExpressionBuilder.scala

@@ -63,11 +63,11 @@ case class CatalystExpressionBuilder(private val function: AnyRef) extends Loggi
    * @param children a sequence of catalyst arguments to the udf.
    * @return the compiled expression, optionally
    */
-  def compile(children: Seq[Expression]): Option[Expression] = {
+  def compile(children: Seq[Expression], objref: Option[Expression] = None): Option[Expression] = {
 
     // create starting state, this will be:
     //   State([children expressions], [empty stack], cond = true, expr = None)
-    val entryState = State.makeStartingState(lambdaReflection, children)
+    val entryState = State.makeStartingState(lambdaReflection, children, objref)
 
     // pick first of the Basic Blocks, and start recursing
     val entryBlock = cfg.basicBlocks.head

udf-compiler/src/main/scala/com/nvidia/spark/udf/Instruction.scala

@@ -18,16 +18,16 @@ package com.nvidia.spark.udf
 
 import CatalystExpressionBuilder.simplify
 import java.nio.charset.Charset
+import javassist.Modifier
 import javassist.bytecode.{CodeIterator, Opcode}
 
 import org.apache.spark.SparkException
 import org.apache.spark.internal.Logging
-import org.apache.spark.sql.catalyst.ScalaReflection
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.types._
 
 
-private object Repr {
+private[udf] object Repr {
 
   abstract class CompilerInternal(name: String) extends Expression {
     override def dataType: DataType = {
@@ -110,6 +110,8 @@ private object Repr {
 
   case class ClassTag[T](classTag: scala.reflect.ClassTag[T])
       extends CompilerInternal("scala.reflect.ClassTag")
+
+  case class UnknownCapturedArg() extends CompilerInternal("unknown captured arg")
 }
 
 /**
@@ -346,7 +348,8 @@ case class Instruction(opcode: Int, operand: Int, instructionStr: String) extend
           (List[Expression], List[Expression])): State = {
     val State(locals, stack, cond, expr) = state
     val method = lambdaReflection.lookupBehavior(operand)
-    val declaringClassName = method.getDeclaringClass.getName
+    val declaringClass = method.getDeclaringClass
+    val declaringClassName = declaringClass.getName
     val paramTypes = method.getParameterTypes
     val (args, rest) = getArgs(stack, paramTypes.length)
     // We don't support arbitrary calls.
@@ -389,9 +392,32 @@ case class Instruction(opcode: Int, operand: Int, instructionStr: String) extend
     } else if (declaringClassName.equals("java.time.LocalDateTime")) {
       State(locals, localDateTimeOp(method.getName, args) :: rest, cond, expr)
     } else {
-      // Other functions
-      throw new SparkException(
-        s"Unsupported instruction: ${Opcode.INVOKEVIRTUAL} ${declaringClassName}")
+      val mModifiers = method.getModifiers
+      val cModifiers = declaringClass.getModifiers
+      if (!javassist.Modifier.isEnum(mModifiers) &&
+          !javassist.Modifier.isInterface(mModifiers) &&
+          !javassist.Modifier.isNative(mModifiers) &&
+          !javassist.Modifier.isPackage(mModifiers) &&
+          !javassist.Modifier.isStrict(mModifiers) &&
+          !javassist.Modifier.isSynchronized(mModifiers) &&
+          !javassist.Modifier.isTransient(mModifiers) &&
+          !javassist.Modifier.isVarArgs(mModifiers) &&
+          !javassist.Modifier.isVolatile(mModifiers) &&
+          (javassist.Modifier.isFinal(mModifiers) ||
+           javassist.Modifier.isFinal(cModifiers))) {
+        val retval = {
+          if (javassist.Modifier.isStatic(mModifiers)) {
+            CatalystExpressionBuilder(method).compile(args)
+          } else {
+            CatalystExpressionBuilder(method).compile(args.tail, Some(args.head))
+          }
+        }
+        State(locals, retval.toList ::: rest, cond, expr)
+      } else {
+        // Other functions
+        throw new SparkException(
+          s"Unsupported invocation of ${declaringClassName}.${method.getName}")
+      }
     }
   }
 

udf-compiler/src/main/scala/com/nvidia/spark/udf/LambdaReflection.scala

@@ -16,12 +16,14 @@
 
 package com.nvidia.spark.udf
 
+import Repr.UnknownCapturedArg
 import java.lang.invoke.SerializedLambda
-
-import javassist.{ClassClassPath, ClassPool, CtBehavior, CtClass, CtField}
-import javassist.bytecode.{CodeIterator, ConstPool, Descriptor}
+import javassist.{ClassClassPath, ClassPool, CtBehavior, CtClass, CtField, CtMethod}
+import javassist.bytecode.{AccessFlag, CodeIterator, ConstPool,
+                           Descriptor, MethodParametersAttribute}
 
 import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.Expression
 import org.apache.spark.sql.types._
 
 //
@@ -30,8 +32,10 @@ import org.apache.spark.sql.types._
 // Provides the interface the class and the method that implements the body of the lambda
 // used by the rest of the compiler.
 //
-case class LambdaReflection(private val classPool: ClassPool,
-    private val serializedLambda: SerializedLambda) {
+class LambdaReflection private(private val classPool: ClassPool,
+                               private val ctClass: CtClass,
+                               private val ctMethod: CtMethod,
+                               val capturedArgs: Seq[Expression] = Seq()) {
   def lookupConstant(constPoolIndex: Int): Any = {
     constPool.getTag(constPoolIndex) match {
       case ConstPool.CONST_Integer => constPool.getIntegerInfo(constPoolIndex)
@@ -61,10 +65,10 @@ case class LambdaReflection(private val classPool: ClassPool,
     val methodName = constPool.getMethodrefName(constPoolIndex)
     val descriptor = constPool.getMethodrefType(constPoolIndex)
     val className = constPool.getMethodrefClassName(constPoolIndex)
-    val params = Descriptor.getParameterTypes(descriptor, classPool)
     if (constPool.isConstructor(className, constPoolIndex) == 0) {
-      classPool.getCtClass(className).getDeclaredMethod(methodName, params)
+      classPool.getCtClass(className).getMethod(methodName, descriptor)
     } else {
+      val params = Descriptor.getParameterTypes(descriptor, classPool)
       classPool.getCtClass(className).getDeclaredConstructor(params)
     }
   }
@@ -76,20 +80,6 @@ case class LambdaReflection(private val classPool: ClassPool,
     constPool.getClassInfo(constPoolIndex)
   }
 
-  // Get the CtClass object for the class that capture the lambda.
-  private val ctClass = {
-    val name = serializedLambda.getCapturingClass.replace('/', '.')
-    val classForName = LambdaReflection.getClass(name)
-    classPool.insertClassPath(new ClassClassPath(classForName))
-    classPool.getCtClass(name)
-  }
-
-  // Get the CtMethod object for the method that implements the lambda body.
-  private val ctMethod = {
-    val lambdaImplName = serializedLambda.getImplMethodName
-    ctClass.getDeclaredMethod(lambdaImplName.stripSuffix("$adapted"))
-  }
-
   private val methodInfo = ctMethod.getMethodInfo
 
   val constPool = methodInfo.getConstPool
@@ -107,6 +97,9 @@ case class LambdaReflection(private val classPool: ClassPool,
 
 object LambdaReflection {
   def apply(function: AnyRef): LambdaReflection = {
+    if (function.isInstanceOf[CtMethod]) {
+      return LambdaReflection(function.asInstanceOf[CtMethod])
+    }
     // writeReplace is supposed to return an object of SerializedLambda from
     // the function class (See
     // https://docs.oracle.com/javase/8/docs/api/java/lang/invoke/SerializedLambda.html).
@@ -117,9 +110,35 @@ object LambdaReflection {
     writeReplace.setAccessible(true)
     val serializedLambda = writeReplace.invoke(function)
         .asInstanceOf[SerializedLambda]
+    val capturedArgs = Seq.tabulate(serializedLambda.getCapturedArgCount){
+      // Add UnknownCapturedArg expressions to the list of captured arguments
+      // until we figure out how to handle captured variables in various
+      // situations.
+      // This, however, lets us pass captured arguments as long as they are not
+      // evaluated.
+      _ => Repr.UnknownCapturedArg()
+    }
 
     val classPool = new ClassPool(true)
-    LambdaReflection(classPool, serializedLambda)
+    // Get the CtClass object for the class that capture the lambda.
+    val ctClass = {
+      val name = serializedLambda.getCapturingClass.replace('/', '.')
+      val classForName = LambdaReflection.getClass(name)
+      classPool.insertClassPath(new ClassClassPath(classForName))
+      classPool.getCtClass(name)
+    }
+    // Get the CtMethod object for the method that implements the lambda body.
+    val ctMethod = {
+      val lambdaImplName = serializedLambda.getImplMethodName
+      ctClass.getDeclaredMethod(lambdaImplName.stripSuffix("$adapted"))
+    }
+    new LambdaReflection(classPool, ctClass, ctMethod, capturedArgs)
+  }
+
+  private def apply(ctMethod: CtMethod): LambdaReflection = {
+    val ctClass = ctMethod.getDeclaringClass
+    val classPool = ctClass.getClassPool
+    new LambdaReflection(classPool, ctClass, ctMethod)
   }
 
   def getClass(name: String): Class[_] = {

udf-compiler/src/main/scala/com/nvidia/spark/udf/State.scala

@@ -19,6 +19,7 @@ package com.nvidia.spark.udf
 import CatalystExpressionBuilder.simplify
 import javassist.CtClass
 
+import org.apache.spark.SparkException
 import org.apache.spark.sql.catalyst.expressions.{Expression, If, Literal, Or}
 
 /**
@@ -74,7 +75,7 @@ import org.apache.spark.sql.catalyst.expressions.{Expression, If, Literal, Or}
  * @param cond
  * @param expr
  */
-case class State(locals: Array[Expression],
+case class State(locals: IndexedSeq[Expression],
     stack: List[Expression] = List(),
     cond: Expression = Literal.TrueLiteral,
     expr: Option[Expression] = None) {
@@ -117,23 +118,23 @@ case class State(locals: Array[Expression],
 
 object State {
   def makeStartingState(lambdaReflection: LambdaReflection,
-      children: Seq[Expression]): State = {
+                        children: Seq[Expression],
+                        objref: Option[Expression]): State = {
     val max = lambdaReflection.maxLocals
-    val params: Seq[(CtClass, Expression)] = lambdaReflection.parameters.view.zip(children)
-    val (locals, _) = params.foldLeft((new Array[Expression](max), 0)) { (l, p) =>
-      val (locals: Array[Expression], index: Int) = l
-      val (param: CtClass, argExp: Expression) = p
-
-      val newIndex = if (param == CtClass.doubleType || param == CtClass.longType) {
+    val args = lambdaReflection.capturedArgs ++ children
+    val paramTypesAndArgs: Seq[(CtClass, Expression)] = lambdaReflection.parameters.view.zip(args)
+    val locals = paramTypesAndArgs.foldLeft(objref.toVector) { (l, p) =>
+      val (paramType : CtClass, argExp) = p
+      if (paramType == CtClass.doubleType || paramType == CtClass.longType) {
         // Long and Double occupies two slots in the local variable array.
+        // Append null to occupy an extra slot.
         // See https://docs.oracle.com/javase/specs/jvms/se8/html/jvms-2.html#jvms-2.6.1
-        index + 2
+        l :+ argExp :+ null
       } else {
-        index + 1
+        l :+ argExp
       }
-
-      (locals.updated(index, argExp), newIndex)
     }
-    State(locals)
+    // Ensure locals have enough slots with padTo.
+    State(locals.padTo(max, null))
   }
 }