[SPARK-26129][SQL] Instrumentation for per-query planning time

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/QueryPlanningTracker.scala

@@ -0,0 +1,109 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst
+
+import scala.collection.JavaConverters._
+
+import org.apache.spark.util.BoundedPriorityQueue
+
+
+/**
+ * A simple utility for tracking runtime and associated stats in query planning.
+ *
+ * There are two separate concepts we track:
+ *
+ * 1. Phases: These are broad scope phases in query planning, as listed below, i.e. analysis,
+ * optimizationm and physical planning (just planning).
+ *
+ * 2. Rules: These are the individual Catalyst rules that we track. In addition to time, we also
+ * track the number of invocations and effective invocations.
+ */
+object QueryPlanningTracker {
+
+  // Define a list of common phases here.
+  val PARSING = "parsing"
+  val ANALYSIS = "analysis"
+  val OPTIMIZATION = "optimization"
+  val PLANNING = "planning"
+
+  class RuleSummary(
+    var totalTimeNs: Long, var numInvocations: Long, var numEffectiveInvocations: Long) {
+
+    def this() = this(totalTimeNs = 0, numInvocations = 0, numEffectiveInvocations = 0)
+
+    override def toString: String = {
+      s"RuleSummary($totalTimeNs, $numInvocations, $numEffectiveInvocations)"
+    }
+  }
+}
+
+
+class QueryPlanningTracker {
+
+  import QueryPlanningTracker._
+
+  // Mapping from the name of a rule to a rule's summary.
+  // Use a Java HashMap for less overhead.
+  private val rulesMap = new java.util.HashMap[String, RuleSummary]
+
+  // From a phase to time in ns.
+  private val phaseToTimeNs = new java.util.HashMap[String, Long]
+
+  /** Measure the runtime of function f, and add it to the time for the specified phase. */
+  def measureTime[T](phase: String)(f: => T): T = {
+    val startTime = System.nanoTime()
+    val ret = f
+    val timeTaken = System.nanoTime() - startTime
+    phaseToTimeNs.put(phase, phaseToTimeNs.getOrDefault(phase, 0) + timeTaken)
+    ret
+  }
+
+  /**
+   * Reecord a specific invocation of a rule.
+   *
+   * @param rule name of the rule
+   * @param timeNs time taken to run this invocation
+   * @param effective whether the invocation has resulted in a plan change
+   */
+  def recordRuleInvocation(rule: String, timeNs: Long, effective: Boolean): Unit = {
+    var s = rulesMap.get(rule)
+    if (s eq null) {
+      s = new RuleSummary
+      rulesMap.put(rule, s)
+    }
+
+    s.totalTimeNs += timeNs
+    s.numInvocations += 1
+    s.numEffectiveInvocations += (if (effective) 1 else 0)
+  }
+
+  // ------------ reporting functions below ------------
+
+  def rules: Map[String, RuleSummary] = rulesMap.asScala.toMap
+
+  def phases: Map[String, Long] = phaseToTimeNs.asScala.toMap
+
+  /** Returns the top k most expensive rules (as measured by time). */
+  def topRulesByTime(k: Int): Seq[(String, RuleSummary)] = {
+    val orderingByTime: Ordering[(String, RuleSummary)] = Ordering.by(e => e._2.totalTimeNs)
+    val q = new BoundedPriorityQueue(k)(orderingByTime)
+    rulesMap.asScala.foreach(q.+=)
+    q.toSeq.sortBy(r => -r._2.totalTimeNs)
+  }
+
+}

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala

@@ -102,29 +102,34 @@ class Analyzer(
     this(catalog, conf, conf.optimizerMaxIterations)
   }
 
-  def executeAndCheck(plan: LogicalPlan): LogicalPlan = AnalysisHelper.markInAnalyzer {
-    val analyzed = execute(plan)
-    try {
-      checkAnalysis(analyzed)
-      analyzed
-    } catch {
-      case e: AnalysisException =>
-        val ae = new AnalysisException(e.message, e.line, e.startPosition, Option(analyzed))
-        ae.setStackTrace(e.getStackTrace)
-        throw ae
+  def executeAndCheck(plan: LogicalPlan, tracker: Option[QueryPlanningTracker]): LogicalPlan = {
+    AnalysisHelper.markInAnalyzer {
+      val analyzed = execute(plan, tracker)
+      try {
+        checkAnalysis(analyzed)
+        analyzed
+      } catch {
+        case e: AnalysisException =>
+          val ae = new AnalysisException(e.message, e.line, e.startPosition, Option(analyzed))
+          ae.setStackTrace(e.getStackTrace)
+          throw ae
+      }
     }
   }
 
-  override def execute(plan: LogicalPlan): LogicalPlan = {
+  override def execute(plan: LogicalPlan, tracker: Option[QueryPlanningTracker]): LogicalPlan = {
     AnalysisContext.reset()
     try {
-      executeSameContext(plan)
+      executeSameContext(plan, tracker)
     } finally {
       AnalysisContext.reset()
     }
   }
 
-  private def executeSameContext(plan: LogicalPlan): LogicalPlan = super.execute(plan)
+  private def executeSameContext(
+      plan: LogicalPlan, tracker: Option[QueryPlanningTracker]): LogicalPlan = {
+    super.execute(plan, tracker)
+  }
 
   def resolver: Resolver = conf.resolver
 
@@ -211,7 +216,7 @@ class Analyzer(
       case With(child, relations) =>
         substituteCTE(child, relations.foldLeft(Seq.empty[(String, LogicalPlan)]) {
           case (resolved, (name, relation)) =>
-            resolved :+ name -> executeSameContext(substituteCTE(relation, resolved))
+            resolved :+ name -> executeSameContext(substituteCTE(relation, resolved), None)
         })
       case other => other
     }
@@ -696,7 +701,7 @@ class Analyzer(
               s"avoid errors. Increase the value of ${SQLConf.MAX_NESTED_VIEW_DEPTH.key} to work " +
               "around this.")
           }
-          executeSameContext(child)
+          executeSameContext(child, None)
         }
         view.copy(child = newChild)
       case p @ SubqueryAlias(_, view: View) =>
@@ -1405,7 +1410,7 @@ class Analyzer(
       do {
         // Try to resolve the subquery plan using the regular analyzer.
         previous = current
-        current = executeSameContext(current)
+        current = executeSameContext(current, None)
 
         // Use the outer references to resolve the subquery plan if it isn't resolved yet.
         val i = plans.iterator
@@ -1527,7 +1532,7 @@ class Analyzer(
               grouping,
               Alias(cond, "havingCondition")() :: Nil,
               child)
-          val resolvedOperator = executeSameContext(aggregatedCondition)
+          val resolvedOperator = executeSameContext(aggregatedCondition, None)
           def resolvedAggregateFilter =
             resolvedOperator
               .asInstanceOf[Aggregate]
@@ -1588,7 +1593,7 @@ class Analyzer(
             unresolvedSortOrders.map(o => Alias(o.child, "aggOrder")())
           val aggregatedOrdering = aggregate.copy(aggregateExpressions = aliasedOrdering)
           val resolvedAggregate: Aggregate =
-            executeSameContext(aggregatedOrdering).asInstanceOf[Aggregate]
+            executeSameContext(aggregatedOrdering, None).asInstanceOf[Aggregate]
           val resolvedAliasedOrdering: Seq[Alias] =
             resolvedAggregate.aggregateExpressions.asInstanceOf[Seq[Alias]]
 

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/rules/RuleExecutor.scala

@@ -18,6 +18,7 @@
 package org.apache.spark.sql.catalyst.rules
 
 import org.apache.spark.internal.Logging
+import org.apache.spark.sql.catalyst.QueryPlanningTracker
 import org.apache.spark.sql.catalyst.errors.TreeNodeException
 import org.apache.spark.sql.catalyst.trees.TreeNode
 import org.apache.spark.sql.catalyst.util.sideBySide
@@ -66,11 +67,14 @@ abstract class RuleExecutor[TreeType <: TreeNode[_]] extends Logging {
    */
   protected def isPlanIntegral(plan: TreeType): Boolean = true
 
+  /** A special overload (as opposed to default parameter values) to allow execute in a closure. */
+  def execute(plan: TreeType): TreeType = execute(plan, None)
+
   /**
    * Executes the batches of rules defined by the subclass. The batches are executed serially
    * using the defined execution strategy. Within each batch, rules are also executed serially.
    */
-  def execute(plan: TreeType): TreeType = {
+  def execute(plan: TreeType, tracker: Option[QueryPlanningTracker]): TreeType = {
     var curPlan = plan
     val queryExecutionMetrics = RuleExecutor.queryExecutionMeter
     val planChangeLogger = new PlanChangeLogger()
@@ -88,15 +92,18 @@ abstract class RuleExecutor[TreeType <: TreeNode[_]] extends Logging {
             val startTime = System.nanoTime()
             val result = rule(plan)
             val runTime = System.nanoTime() - startTime
+            val effective = !result.fastEquals(plan)
 
-            if (!result.fastEquals(plan)) {
+            if (effective) {
               queryExecutionMetrics.incNumEffectiveExecution(rule.ruleName)
               queryExecutionMetrics.incTimeEffectiveExecutionBy(rule.ruleName, runTime)
               planChangeLogger.log(rule.ruleName, plan, result)
             }
             queryExecutionMetrics.incExecutionTimeBy(rule.ruleName, runTime)
             queryExecutionMetrics.incNumExecution(rule.ruleName)
 
+            tracker.foreach(_.recordRuleInvocation(rule.ruleName, runTime, effective))
+
             // Run the structural integrity checker against the plan after each rule.
             if (!isPlanIntegral(result)) {
               val message = s"After applying rule ${rule.ruleName} in batch ${batch.name}, " +

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/QueryPlanningTrackerSuite.scala

@@ -0,0 +1,78 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst
+
+import org.apache.spark.SparkFunSuite
+
+class QueryPlanningTrackerSuite extends SparkFunSuite {
+
+  test("phases") {
+    val t = new QueryPlanningTracker
+    t.measureTime("p1") {
+      Thread.sleep(1)
+    }
+
+    assert(t.phases("p1") > 0)
+    assert(!t.phases.contains("p2"))
+
+    val old = t.phases("p1")
+
+    t.measureTime("p1") {
+      Thread.sleep(1)
+    }
+    assert(t.phases("p1") > old)
+  }
+
+  test("rules") {
+    val t = new QueryPlanningTracker
+    t.recordRuleInvocation("r1", 1, effective = false)
+    t.recordRuleInvocation("r2", 2, effective = true)
+    t.recordRuleInvocation("r3", 1, effective = false)
+    t.recordRuleInvocation("r3", 2, effective = true)
+
+    val rules = t.rules
+
+    assert(rules("r1").totalTimeNs == 1)
+    assert(rules("r1").numInvocations == 1)
+    assert(rules("r1").numEffectiveInvocations == 0)
+
+    assert(rules("r2").totalTimeNs == 2)
+    assert(rules("r2").numInvocations == 1)
+    assert(rules("r2").numEffectiveInvocations == 1)
+
+    assert(rules("r3").totalTimeNs == 3)
+    assert(rules("r3").numInvocations == 2)
+    assert(rules("r3").numEffectiveInvocations == 1)
+  }
+
+  test("topRulesByTime") {
+    val t = new QueryPlanningTracker
+    t.recordRuleInvocation("r2", 2, effective = true)
+    t.recordRuleInvocation("r4", 4, effective = true)
+    t.recordRuleInvocation("r1", 1, effective = false)
+    t.recordRuleInvocation("r3", 3, effective = false)
+
+    val top = t.topRulesByTime(2)
+    assert(top.size == 2)
+    assert(top(0)._1 == "r4")
+    assert(top(1)._1 == "r3")
+
+    // Don't crash when k > total size
+    assert(t.topRulesByTime(10).size == 4)
+  }
+}

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisTest.scala

@@ -54,7 +54,7 @@ trait AnalysisTest extends PlanTest {
       expectedPlan: LogicalPlan,
       caseSensitive: Boolean = true): Unit = {
     val analyzer = getAnalyzer(caseSensitive)
-    val actualPlan = analyzer.executeAndCheck(inputPlan)
+    val actualPlan = analyzer.executeAndCheck(inputPlan, None)
     comparePlans(actualPlan, expectedPlan)
   }
 

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ResolveGroupingAnalyticsSuite.scala

@@ -109,7 +109,7 @@ class ResolveGroupingAnalyticsSuite extends AnalysisTest {
       Seq(UnresolvedAlias(Multiply(unresolved_a, Literal(2))),
         unresolved_b, UnresolvedAlias(count(unresolved_c))))
 
-    val resultPlan = getAnalyzer(true).executeAndCheck(originalPlan2)
+    val resultPlan = getAnalyzer(true).executeAndCheck(originalPlan2, None)
     val gExpressions = resultPlan.asInstanceOf[Aggregate].groupingExpressions
     assert(gExpressions.size == 3)
     val firstGroupingExprAttrName =

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ResolvedUuidExpressionsSuite.scala

@@ -47,7 +47,7 @@ class ResolvedUuidExpressionsSuite extends AnalysisTest {
 
   test("analyzed plan sets random seed for Uuid expression") {
     val plan = r.select(a, uuid1)
-    val resolvedPlan = analyzer.executeAndCheck(plan)
+    val resolvedPlan = analyzer.executeAndCheck(plan, None)
     getUuidExpressions(resolvedPlan).foreach { u =>
       assert(u.resolved)
       assert(u.randomSeed.isDefined)
@@ -56,14 +56,14 @@ class ResolvedUuidExpressionsSuite extends AnalysisTest {
 
   test("Uuid expressions should have different random seeds") {
     val plan = r.select(a, uuid1).groupBy(uuid1Ref)(uuid2, uuid3)
-    val resolvedPlan = analyzer.executeAndCheck(plan)
+    val resolvedPlan = analyzer.executeAndCheck(plan, None)
     assert(getUuidExpressions(resolvedPlan).map(_.randomSeed.get).distinct.length == 3)
   }
 
   test("Different analyzed plans should have different random seeds in Uuids") {
     val plan = r.select(a, uuid1).groupBy(uuid1Ref)(uuid2, uuid3)
-    val resolvedPlan1 = analyzer.executeAndCheck(plan)
-    val resolvedPlan2 = analyzer.executeAndCheck(plan)
+    val resolvedPlan1 = analyzer.executeAndCheck(plan, None)
+    val resolvedPlan2 = analyzer.executeAndCheck(plan, None)
     val uuids1 = getUuidExpressions(resolvedPlan1)
     val uuids2 = getUuidExpressions(resolvedPlan2)
     assert(uuids1.distinct.length == 3)

sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala

@@ -32,6 +32,7 @@ import org.apache.spark.api.java.function._
 import org.apache.spark.api.python.{PythonRDD, SerDeUtil}
 import org.apache.spark.broadcast.Broadcast
 import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.catalyst.QueryPlanningTracker
 import org.apache.spark.sql.catalyst.analysis._
 import org.apache.spark.sql.catalyst.catalog.HiveTableRelation
 import org.apache.spark.sql.catalyst.encoders._
@@ -76,6 +77,14 @@ private[sql] object Dataset {
     qe.assertAnalyzed()
     new Dataset[Row](sparkSession, qe, RowEncoder(qe.analyzed.schema))
   }
+
+  /** A variant of ofRows that allows passing in a tracker so we can track query parsing time. */
+  def ofRows(sparkSession: SparkSession, logicalPlan: LogicalPlan, tracker: QueryPlanningTracker)
+    : DataFrame = {
+    val qe = new QueryExecution(sparkSession, logicalPlan, tracker)
+    qe.assertAnalyzed()
+    new Dataset[Row](sparkSession, qe, RowEncoder(qe.analyzed.schema))
+  }
 }
 
 /**