Address reviews

Author: maropu

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

@@ -1214,15 +1214,11 @@ abstract class CodeGenerator[InType <: AnyRef, OutType <: AnyRef] extends Loggin
 /**
  * Java bytecode statistics of a compiled class by Janino.
  */
-case class ByteCodeStats(maxClassCodeSize: Int, maxMethodCodeSize: Int, maxConstPoolSize: Int)
+case class ByteCodeStats(
+  maxClassCodeSize: Int, maxMethodCodeSize: Int, maxConstPoolSize: Int, numInnerClasses: Int)
 
 object ByteCodeStats {
-
-  val unavailable = ByteCodeStats(-1, -1, -1)
-
-  def apply(codeStats: (Int, Int, Int)): ByteCodeStats = {
-    ByteCodeStats(codeStats._1, codeStats._2, codeStats._3)
-  }
+  val UNAVAILABLE = ByteCodeStats(-1, -1, -1, -1)
 }
 
 object CodeGenerator extends Logging {
@@ -1332,9 +1328,9 @@ object CodeGenerator extends Logging {
   }
 
   /**
-   * Returns the bytecode statistics (max class bytecode size, max method bytecode size, and
-   * max constant pool size) of generated classes by inspecting janino private fields.inspecting
-   * janino private fields. Also, this method updates the metrics information.
+   * Returns the bytecode statistics (max class bytecode size, max method bytecode size,
+   * max constant pool size, and # of inner classes) of generated classes
+   * by inspecting Janino classes. Also, this method updates the metrics information.
    */
   private def updateAndGetCompilationStats(evaluator: ClassBodyEvaluator): ByteCodeStats = {
     // First retrieve the generated classes.
@@ -1378,9 +1374,17 @@ object CodeGenerator extends Logging {
       }
     }
 
-    ByteCodeStats(codeStats.reduce[(Int, Int, Int)] { case (v1, v2) =>
+    // Computes the max values of the three metrics: class code size, method code size,
+    // and constant pool size.
+    val maxCodeStats = codeStats.reduce[(Int, Int, Int)] { case (v1, v2) =>
       (Math.max(v1._1, v2._1), Math.max(v1._2, v2._2), Math.max(v1._3, v2._3))
-    })
+    }
+    ByteCodeStats(
+      maxClassCodeSize = maxCodeStats._1,
+      maxMethodCodeSize = maxCodeStats._2,
+      maxConstPoolSize = maxCodeStats._3,
+      // Minus 2 for `GeneratedClass` and an outer-most generated class
+      numInnerClasses = classes.size - 2)
   }
 
   /**

sql/core/src/main/scala/org/apache/spark/sql/execution/debug/package.scala

@@ -83,9 +83,10 @@ package object debug {
     val codegenSeq = codegenStringSeq(plan)
     append(s"Found ${codegenSeq.size} WholeStageCodegen subtrees.\n")
     for (((subtree, code, codeStats), i) <- codegenSeq.zipWithIndex) {
-      val codeStatsStr = s"maxClassCodeSize:${codeStats.maxClassCodeSize} " +
-        s"maxMethodCodeSize:${codeStats.maxMethodCodeSize} " +
-        s"maxConstantPoolSize:${codeStats.maxConstPoolSize}"
+      val codeStatsStr = s"maxClassCodeSize:${codeStats.maxClassCodeSize}; " +
+        s"maxMethodCodeSize:${codeStats.maxMethodCodeSize}; " +
+        s"maxConstantPoolSize:${codeStats.maxConstPoolSize}; " +
+        s"numInnerClasses:${codeStats.numInnerClasses}"
       append(s"== Subtree ${i + 1} / ${codegenSeq.size} ($codeStatsStr) ==\n")
       append(subtree)
       append("\nGenerated code:\n")
@@ -113,7 +114,7 @@ package object debug {
         CodeGenerator.compile(source)._2
       } catch {
         case NonFatal(_) =>
-          ByteCodeStats.unavailable
+          ByteCodeStats.UNAVAILABLE
       }
       (subtree.toString, CodeFormatter.format(source), codeStats)
     }

sql/core/src/test/scala/org/apache/spark/sql/execution/WholeStageCodegenSuite.scala

@@ -213,10 +213,10 @@ class WholeStageCodegenSuite extends QueryTest with SharedSparkSession {
 
   ignore("SPARK-21871 check if we can get large code size when compiling too long functions") {
     val codeWithShortFunctions = genGroupByCode(3)
-    val (_, ByteCodeStats(_, maxCodeSize1, _)) = CodeGenerator.compile(codeWithShortFunctions)
+    val (_, ByteCodeStats(_, maxCodeSize1, _, _)) = CodeGenerator.compile(codeWithShortFunctions)
     assert(maxCodeSize1 < SQLConf.WHOLESTAGE_HUGE_METHOD_LIMIT.defaultValue.get)
     val codeWithLongFunctions = genGroupByCode(50)
-    val (_, ByteCodeStats(_, maxCodeSize2, _)) = CodeGenerator.compile(codeWithLongFunctions)
+    val (_, ByteCodeStats(_, maxCodeSize2, _, _)) = CodeGenerator.compile(codeWithLongFunctions)
     assert(maxCodeSize2 > SQLConf.WHOLESTAGE_HUGE_METHOD_LIMIT.defaultValue.get)
   }
 

sql/core/src/test/scala/org/apache/spark/sql/execution/debug/DebuggingSuite.scala

@@ -90,4 +90,30 @@ class DebuggingSuite extends SharedSparkSession {
         | id LongType: {}
         |""".stripMargin))
   }
+
+  test("Prints bytecode statistics in debugCodegen") {
+    Seq(("SELECT sum(v) FROM VALUES(1) t(v)", (0, 0)),
+      // We expect HashAggregate uses an inner class for fast hash maps
+      // in partial aggregates with keys.
+      ("SELECT k, avg(v) FROM VALUES((1, 1)) t(k, v) GROUP BY k", (0, 1)))
+        .foreach { case (query, (expectedNumInnerClasses0, expectedNumInnerClasses1)) =>
+
+      val executedPlan = sql(query).queryExecution.executedPlan
+      val res = codegenStringSeq(executedPlan)
+      assert(res.length == 2)
+      assert(res.forall { case (_, _, codeStats) =>
+        codeStats.maxClassCodeSize > 0 &&
+          codeStats.maxMethodCodeSize > 0 &&
+          codeStats.maxConstPoolSize > 0
+      })
+      assert(res(0)._3.numInnerClasses == expectedNumInnerClasses0)
+      assert(res(1)._3.numInnerClasses == expectedNumInnerClasses1)
+
+      val debugCodegenStr = codegenString(executedPlan)
+      assert(debugCodegenStr.contains("maxClassCodeSize:"))
+      assert(debugCodegenStr.contains("maxMethodCodeSize:"))
+      assert(debugCodegenStr.contains("maxConstantPoolSize:"))
+      assert(debugCodegenStr.contains("numInnerClasses:"))
+    }
+  }
 }