add test cases

rebase master

Author: kiszk

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

@@ -1237,8 +1237,9 @@ object CodeGenerator extends Logging {
   // bytecode instruction
   final val MUTABLESTATEARRAY_SIZE_LIMIT = 32768
 
-  private lazy val compilerImpl: CompilerBase = {
-    val compiler = SparkEnv.get.conf.get("spark.sql.javaCompiler", "jdk").toLowerCase(Locale.ROOT)
+  // Make this non-private and non-lazy method for testing
+  def _compilerImpl(): CompilerBase = {
+    val compiler = SQLConf.get.javaCompiler
     val compilerInstance = compiler match {
       case "janino" => JaninoCompiler
       case "jdk" => if (JdkCompiler.javaCompiler != null) JdkCompiler else JaninoCompiler
@@ -1249,12 +1250,14 @@ object CodeGenerator extends Logging {
     compilerInstance
   }
 
+  private lazy val compilerImpl = _compilerImpl
+
   lazy val janinoCompilerEnabled: Boolean = {
     compilerImpl == JaninoCompiler
   }
 
   /**
-   * Compile the Java source code into a Java class, using Janino.
+   * Compile the Java source code into a Java class, using Janino or javac
    *
    * @return a pair of a generated class and the max bytecode size of generated functions.
    */
@@ -1269,7 +1272,7 @@ object CodeGenerator extends Logging {
   }
 
   /**
-   * Compile the Java source code into a Java class, using Janino.
+   * Compile the Java source code into a Java class, using Janino or javac
    */
   private[this] def doCompile(code: CodeAndComment): (GeneratedClass, Int) = {
     compilerImpl.compile(code)

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

@@ -17,18 +17,29 @@
 
 package org.apache.spark.sql.catalyst.expressions.codegen.compiler
 
+import java.io.ByteArrayInputStream
+
+import scala.collection.JavaConverters._
+import scala.language.existentials
+import scala.util.control.NonFatal
+
+import org.codehaus.janino.util.ClassFile
+
 import org.apache.spark.{TaskContext, TaskKilledException}
 import org.apache.spark.executor.InputMetrics
+import org.apache.spark.internal.Logging
+import org.apache.spark.metrics.source.CodegenMetrics
 import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.expressions.{Expression, UnsafeArrayData, UnsafeMapData, UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.codegen.{CodeAndComment, GeneratedClass}
 import org.apache.spark.sql.catalyst.util.{ArrayData, MapData}
 import org.apache.spark.sql.types.Decimal
 import org.apache.spark.unsafe.Platform
 import org.apache.spark.unsafe.types.{CalendarInterval, UTF8String}
+import org.apache.spark.util.Utils
 
 
-abstract class CompilerBase {
+abstract class CompilerBase extends Logging {
   protected val className = "org.apache.spark.sql.catalyst.expressions.GeneratedClass"
 
   protected val importClassNames = Seq(
@@ -77,4 +88,35 @@ abstract class CompilerBase {
   }
 
   def compile(code: CodeAndComment): (GeneratedClass, Int)
+
+  /**
+   * Returns the max bytecode size of the generated functions by inspecting janino private fields.
+   * Also, this method updates the metrics information.
+   */
+  protected def updateAndGetBytecodeSize(byteCodes: Iterable[Array[Byte]]): Int = {
+    // Walk the classes to get at the method bytecode.
+    val codeAttr = Utils.classForName("org.codehaus.janino.util.ClassFile$CodeAttribute")
+    val codeAttrField = codeAttr.getDeclaredField("code")
+    codeAttrField.setAccessible(true)
+    val codeSizes = byteCodes.flatMap { byteCode =>
+      CodegenMetrics.METRIC_GENERATED_CLASS_BYTECODE_SIZE.update(byteCode.size)
+      try {
+        val cf = new ClassFile(new ByteArrayInputStream(byteCode))
+        val stats = cf.methodInfos.asScala.flatMap { method =>
+          method.getAttributes().filter(_.getClass.getName == codeAttr.getName).map { a =>
+            val byteCodeSize = codeAttrField.get(a).asInstanceOf[Array[Byte]].length
+            CodegenMetrics.METRIC_GENERATED_METHOD_BYTECODE_SIZE.update(byteCodeSize)
+            byteCodeSize
+          }
+        }
+        Some(stats)
+      } catch {
+        case NonFatal(e) =>
+          logWarning("Error calculating stats of compiled class.", e)
+          None
+      }
+    }.flatten
+
+    codeSizes.max
+  }
 }

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

@@ -17,30 +17,20 @@
 
 package org.apache.spark.sql.catalyst.expressions.codegen.compiler
 
-import java.io.ByteArrayInputStream
 import java.util.{Map => JavaMap}
 
 import scala.collection.JavaConverters._
-import scala.language.existentials
-import scala.util.control.NonFatal
 
 import org.codehaus.commons.compiler.CompileException
-import org.codehaus.janino.{ByteArrayClassLoader, ClassBodyEvaluator, InternalCompilerException, SimpleCompiler}
-import org.codehaus.janino.util.ClassFile
+import org.codehaus.janino._
 
-import org.apache.spark.internal.Logging
-import org.apache.spark.metrics.source.CodegenMetrics
-import org.apache.spark.sql.catalyst.expressions.codegen.{CodeAndComment, CodeFormatter, GeneratedClass}
+import org.apache.spark.sql.catalyst.expressions.codegen._
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.util.{ParentClassLoader, Utils}
 
 
-object JaninoCompiler extends CompilerBase with Logging {
+object JaninoCompiler extends CompilerBase {
 
-  /**
-   * Returns the max bytecode size of the generated functions by inspecting janino private fields.
-   * Also, this method updates the metrics information.
-   */
   private def updateAndGetCompilationStats(evaluator: ClassBodyEvaluator): Int = {
     // First retrieve the generated classes.
     val classes = {
@@ -51,31 +41,7 @@ object JaninoCompiler extends CompilerBase with Logging {
       classesField.setAccessible(true)
       classesField.get(loader).asInstanceOf[JavaMap[String, Array[Byte]]].asScala
     }
-
-    // Then walk the classes to get at the method bytecode.
-    val codeAttr = Utils.classForName("org.codehaus.janino.util.ClassFile$CodeAttribute")
-    val codeAttrField = codeAttr.getDeclaredField("code")
-    codeAttrField.setAccessible(true)
-    val codeSizes = classes.flatMap { case (_, classBytes) =>
-      CodegenMetrics.METRIC_GENERATED_CLASS_BYTECODE_SIZE.update(classBytes.length)
-      try {
-        val cf = new ClassFile(new ByteArrayInputStream(classBytes))
-        val stats = cf.methodInfos.asScala.flatMap { method =>
-          method.getAttributes().filter(_.getClass.getName == codeAttr.getName).map { a =>
-            val byteCodeSize = codeAttrField.get(a).asInstanceOf[Array[Byte]].length
-            CodegenMetrics.METRIC_GENERATED_METHOD_BYTECODE_SIZE.update(byteCodeSize)
-            byteCodeSize
-          }
-        }
-        Some(stats)
-      } catch {
-        case NonFatal(e) =>
-          logWarning("Error calculating stats of compiled class.", e)
-          None
-      }
-    }.flatten
-
-    codeSizes.max
+    updateAndGetBytecodeSize(classes.values)
   }
 
   override def compile(code: CodeAndComment): (GeneratedClass, Int) = {

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

@@ -25,14 +25,16 @@ import javax.tools._
 import javax.tools.JavaFileManager.Location
 import javax.tools.JavaFileObject.Kind
 
+import scala.collection.JavaConverters._
 import scala.collection.mutable
+import scala.language.existentials
 
 import org.codehaus.commons.compiler.{CompileException, Location => CompilerLocation}
 
 import org.apache.spark.internal.Logging
-import org.apache.spark.metrics.source.CodegenMetrics
 import org.apache.spark.sql.catalyst.expressions.codegen.{CodeAndComment, CodeFormatter, GeneratedClass}
-import org.apache.spark.util.{ParentClassLoader, Utils}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.util.ParentClassLoader
 
 
 class JavaCodeManager(fileManager: JavaFileManager)
@@ -123,13 +125,11 @@ class JDKDiagnosticListener extends DiagnosticListener[JavaFileObject] {
       // Wrap the exception in a RuntimeException, because "report()"
       // does not declare checked exceptions.
       throw new RuntimeException(new CompileException(message, loc))
-    // } else if (logger.isTraceEnabled()) {
-    //   logger.trace(diagnostic.toString() + " (" + diagnostic.getCode() + ")")
     }
   }
 }
 
-object JdkCompiler extends CompilerBase with Logging {
+object JdkCompiler extends CompilerBase {
   val javaCompiler = {
     ToolProvider.getSystemJavaCompiler
   }
@@ -142,7 +142,12 @@ object JdkCompiler extends CompilerBase with Logging {
     if ("-g".equals(debugOption)) {
       debugOption.append("none,")
     }
-    Arrays.asList("-classpath", System.getProperty("java.class.path"), debugOption.init.toString())
+    val compilerOption = SQLConf.get.jdkCompilerOptions
+    if (compilerOption != null) {
+      debugOption.append(compilerOption).append(",")
+    }
+
+    Arrays.asList("-classpath", System.getProperty("java.class.path"), debugOption.toString())
   }
 
   private val listener = new JDKDiagnosticListener()
@@ -194,39 +199,10 @@ object JdkCompiler extends CompilerBase with Logging {
       case _: Throwable => throw new CompileException("Compilation failed", null)
     }
 
-    // TODO: Needs to get the max bytecode size of generated methods
-    val maxMethodBytecodeSize = updateAndGetCompilationStats(fileManager.objects.toMap)
+    val byteCodes = fileManager.objects.toMap.values.map(_.getBytecode)
+    val maxMethodBytecodeSize = updateAndGetBytecodeSize(byteCodes)
 
     val clazz = fileManager.getClassLoader(null).loadClass(clazzName)
     (clazz.newInstance().asInstanceOf[GeneratedClass], maxMethodBytecodeSize)
   }
-
-  private def updateAndGetCompilationStats(objects: Map[String, JavaCode]): Int = {
-    val codeAttr = Utils.classForName("org.codehaus.janino.util.ClassFile$CodeAttribute")
-    val codeAttrField = codeAttr.getDeclaredField("code")
-    codeAttrField.setAccessible(true)
-    /*
-    val codeSizes = objects.foreach { case (_, javaCode) =>
-      CodegenMetrics.METRIC_GENERATED_CLASS_BYTECODE_SIZE.update(javaCode.getBytecode.size)
-      try {
-        val cf = new ClassFile(new ByteArrayInputStream(javaCode.getBytecode))
-        val stats = cf.methodInfos.asScala.flatMap { method =>
-          method.getAttributes().filter(_.getClass.getName == codeAttr.getName).map { a =>
-            val byteCodeSize = codeAttrField.get(a).asInstanceOf[Array[Byte]].length
-            CodegenMetrics.METRIC_GENERATED_METHOD_BYTECODE_SIZE.update(byteCodeSize)
-            byteCodeSize
-          }
-        }
-        Some(stats)
-      } catch {
-        case NonFatal(e) =>
-          logWarning("Error calculating stats of compiled class.", e)
-          None
-      }
-    }.flatten
-
-    codeSizes.max
-    */
-    0
-  }
 }

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala

@@ -1740,6 +1740,10 @@ class SQLConf extends Serializable with Logging {
 
   def codegenCacheMaxEntries: Int = getConf(StaticSQLConf.CODEGEN_CACHE_MAX_ENTRIES)
 
+  def javaCompiler: String = getConf(StaticSQLConf.CODEGEN_JAVA_COMPILER)
+
+  def jdkCompilerOptions: String = getConf(StaticSQLConf.CODEGEN_JDK_JAVA_COMPILER_OPTION)
+
   def exchangeReuseEnabled: Boolean = getConf(EXCHANGE_REUSE_ENABLED)
 
   def caseSensitiveAnalysis: Boolean = getConf(SQLConf.CASE_SENSITIVE)

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/StaticSQLConf.scala

@@ -82,6 +82,22 @@ object StaticSQLConf {
     .booleanConf
     .createWithDefault(false)
 
+  val CODEGEN_JAVA_COMPILER = buildStaticConf("spark.sql.codegen.javaCompiler")
+      .internal()
+      .doc("Sets the Java bytecode compiler for compiling Java methods for DataFrame or Dataset " +
+        "program. Acceptable values include: jdk or janino")
+      .stringConf
+      .checkValues(Set("jdk", "janino"))
+      .createWithDefault("jdk")
+
+  val CODEGEN_JDK_JAVA_COMPILER_OPTION =
+    buildStaticConf("spark.sql.codegen.javaCompiler.jdkOption")
+      .internal()
+      .doc("Sets compiler options for JDK Java bytecode compiler. This is ignored when janino is" +
+        "is selected")
+      .stringConf
+      .createWithDefault("")
+
   // When enabling the debug, Spark SQL internal table properties are not filtered out; however,
   // some related DDL commands (e.g., ANALYZE TABLE and CREATE TABLE LIKE) might not work properly.
   val DEBUG_MODE = buildStaticConf("spark.sql.debug")