Initial implementation of row count estimates in cost-based optimizer

shims/spark301/src/main/scala/com/nvidia/spark/rapids/shims/spark301/GpuBroadcastExchangeExec.scala

@@ -17,6 +17,8 @@ package com.nvidia.spark.rapids.shims.spark301
 
 import java.util.UUID
 
+import com.nvidia.spark.rapids.GpuMetric
+
 import org.apache.spark.sql.catalyst.plans.logical.Statistics
 import org.apache.spark.sql.catalyst.plans.physical.BroadcastMode
 import org.apache.spark.sql.execution.SparkPlan
@@ -30,8 +32,9 @@ case class GpuBroadcastExchangeExec(
   override def runId: UUID = _runId
 
   override def runtimeStatistics: Statistics = {
-    val dataSize = metrics("dataSize").value
-    Statistics(dataSize)
+    Statistics(
+      sizeInBytes = metrics("dataSize").value,
+      rowCount = Some(metrics(GpuMetric.NUM_OUTPUT_ROWS).value))
   }
 
   override def doCanonicalize(): SparkPlan = {

sql-plugin/src/main/scala/com/nvidia/spark/rapids/CostBasedOptimizer.scala

@@ -20,10 +20,12 @@ import scala.collection.mutable.ListBuffer
 
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.catalyst.expressions.{Alias, AttributeReference, Expression}
-import org.apache.spark.sql.execution.{ProjectExec, SparkPlan}
-import org.apache.spark.sql.execution.adaptive.CustomShuffleReaderExec
+import org.apache.spark.sql.catalyst.plans.{JoinType, LeftAnti, LeftSemi}
+import org.apache.spark.sql.execution.{GlobalLimitExec, LocalLimitExec, ProjectExec, SparkPlan, TakeOrderedAndProjectExec, UnionExec}
+import org.apache.spark.sql.execution.adaptive.{CustomShuffleReaderExec, QueryStageExec}
+import org.apache.spark.sql.execution.aggregate.HashAggregateExec
 import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec
-import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, ShuffledHashJoinExec}
+import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, ShuffledHashJoinExec, SortMergeJoinExec}
 import org.apache.spark.sql.internal.SQLConf
 
 class CostBasedOptimizer(conf: RapidsConf) extends Logging {
@@ -78,6 +80,8 @@ class CostBasedOptimizer(conf: RapidsConf) extends Logging {
     val totalCpuCost = operatorCpuCost + childCpuCosts.sum
     var totalGpuCost = operatorGpuCost + childGpuCosts.sum
 
+    plan.estimatedOutputRows = RowCountPlanVisitor.visit(plan)
+
     // determine how many transitions between CPU and GPU are taking place between
     // the child operators and this operator
     val numTransitions = plan.childPlans
@@ -270,6 +274,69 @@ class DefaultCostModel(conf: RapidsConf) extends CostModel {
 
 }
 
+/**
+ * Estimate the number of rows that an operator will output. Note that these row counts are
+ * the aggregate across all output partitions.
+ *
+ * Logic is based on Spark's SizeInBytesOnlyStatsPlanVisitor. which operates on logical plans
+ * and only computes data sizes, not row counts.
+ */
+object RowCountPlanVisitor {
+
+  def visit(plan: SparkPlanMeta[_]): Option[BigInt] = plan.wrapped match {
+    case p: QueryStageExec =>
+      ShimLoader.getSparkShims.getQueryStageRuntimeStatistics(p).rowCount
+    case GlobalLimitExec(limit, _) =>
+      visit(plan.childPlans.head).map(_.min(limit)).orElse(Some(limit))
+    case LocalLimitExec(limit, _) =>
+      // LocalLimit applies the same limit for each partition
+      val n = limit * plan.wrapped.asInstanceOf[SparkPlan]
+          .outputPartitioning.numPartitions
+      visit(plan.childPlans.head).map(_.min(n)).orElse(Some(n))
+    case p: TakeOrderedAndProjectExec =>
+      visit(plan.childPlans.head).map(_.min(p.limit)).orElse(Some(p.limit))
+    case p: HashAggregateExec if p.groupingExpressions.isEmpty =>
+      Some(1)
+    case p: SortMergeJoinExec =>
+      estimateJoin(plan, p.joinType)
+    case p: ShuffledHashJoinExec =>
+      estimateJoin(plan, p.joinType)
+    case p: BroadcastHashJoinExec =>
+      estimateJoin(plan, p.joinType)
+    case _: UnionExec =>
+      Some(plan.childPlans.flatMap(visit).sum)
+    case _ =>
+      default(plan)
+  }
+
+  private def estimateJoin(plan: SparkPlanMeta[_], joinType: JoinType): Option[BigInt] = {
+    joinType match {
+      case LeftAnti | LeftSemi =>
+        // LeftSemi and LeftAnti won't ever be bigger than left
+        visit(plan.childPlans.head)
+      case _ =>
+        default(plan)
+    }
+  }
+
+  /**
+   * The default row count is the product of the row count of all child plans.
+   */
+  private def default(p: SparkPlanMeta[_]): Option[BigInt] = {
+    val one = BigInt(1)
+    val product = p.childPlans.map(visit)
+        .filter(_.exists(_ > 0L))
+        .map(_.get)
+        .product
+    if (product == one) {
+      // product will be 1 when there are no child plans
+      None
+    } else {
+      Some(product)
+    }
+  }
+}
+
 sealed abstract class Optimization
 
 case class AvoidTransition[INPUT <: SparkPlan](plan: SparkPlanMeta[INPUT]) extends Optimization {

sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsMeta.scala

@@ -540,6 +540,7 @@ abstract class SparkPlanMeta[INPUT <: SparkPlan](plan: INPUT,
 
   var cpuCost: Double = 0
   var gpuCost: Double = 0
+  var estimatedOutputRows: Option[BigInt] = None
 
   override def convertToCpu(): SparkPlan = {
     wrapped.withNewChildren(childPlans.map(_.convertIfNeeded()))

tests/src/test/scala/com/nvidia/spark/rapids/AdaptiveQueryExecSuite.scala

@@ -550,25 +550,6 @@ class AdaptiveQueryExecSuite
     }, conf)
   }
 
-  /** most of the AQE tests requires Spark 3.0.1 or later */
-  private def assumeSpark301orLater =
-    assume(cmpSparkVersion(3, 0, 1) >= 0)
-
-  private def assumePriorToSpark320 =
-    assume(cmpSparkVersion(3, 2, 0) < 0)
-
-  private def cmpSparkVersion(major: Int, minor: Int, bugfix: Int): Int = {
-    val sparkShimVersion = ShimLoader.getSparkShims.getSparkShimVersion
-    val (sparkMajor, sparkMinor, sparkBugfix) = sparkShimVersion match {
-      case SparkShimVersion(a, b, c) => (a, b, c)
-      case DatabricksShimVersion(a, b, c) => (a, b, c)
-      case EMRShimVersion(a, b, c) => (a, b, c)
-    }
-    val fullVersion = ((major.toLong * 1000) + minor) * 1000 + bugfix
-    val sparkFullVersion = ((sparkMajor.toLong * 1000) + sparkMinor) * 1000 + sparkBugfix
-    sparkFullVersion.compareTo(fullVersion)
-  }
-
   def checkSkewJoin(
       joins: Seq[GpuShuffledHashJoinBase],
       leftSkewNum: Int,

tests/src/test/scala/com/nvidia/spark/rapids/CostBasedOptimizerSuite.scala

@@ -390,6 +390,111 @@ class CostBasedOptimizerSuite extends SparkQueryCompareTestSuite with BeforeAndA
     }, conf)
   }
 
+  test("Compute estimated row count nested joins no broadcast") {
+    assumeSpark301orLater
+
+    val conf = new SparkConf()
+        .set(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key, "true")
+        .set(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key, "-1")
+        .set(RapidsConf.OPTIMIZER_ENABLED.key, "true")
+        .set(RapidsConf.TEST_ALLOWED_NONGPU.key,
+          "ProjectExec,SortMergeJoinExec,SortExec,Alias,Cast,LessThan")
+
+    var plans: ListBuffer[SparkPlanMeta[SparkPlan]] =
+      new ListBuffer[SparkPlanMeta[SparkPlan]]()
+    GpuOverrides.addListener(
+      (plan: SparkPlanMeta[SparkPlan],
+          _: SparkPlan,
+          _: Seq[Optimization]) => {
+        plans += plan
+      })
+
+    withGpuSparkSession(spark => {
+      val df1: DataFrame = createQuery(spark).alias("l")
+      val df2: DataFrame = createQuery(spark).alias("r")
+      val df = df1.join(df2,
+        col("l.more_strings_1").equalTo(col("r.more_strings_2")))
+      df.collect()
+    }, conf)
+
+    val accum = new ListBuffer[SparkPlanMeta[_]]()
+    plans.foreach(collectPlansWithRowCount(_, accum))
+
+    val summary = accum
+        .map(plan => plan.wrapped.getClass.getSimpleName -> plan.estimatedOutputRows.get)
+        .distinct
+        .sorted
+
+    // we expect to see:
+    // 4 rows for the leaf query stages
+    // 16 rows for the nested joins (4 x 4)
+    // 256 rows for the final join (16 x 16)
+    assert(summary === Seq(
+      "CustomShuffleReaderExec" -> 4,
+      "ProjectExec" -> 4,
+      "ProjectExec" -> 16,
+      "ShuffleExchangeExec" -> 4,
+      "ShuffleExchangeExec" -> 16,
+      "ShuffleQueryStageExec" -> 4,
+      "SortExec" -> 4,
+      "SortExec" -> 16,
+      "SortMergeJoinExec" -> 16,
+      "SortMergeJoinExec" -> 256))
+  }
+
+  test("Compute estimated row count nested joins with broadcast") {
+    assumeSpark301orLater
+
+    val conf = new SparkConf()
+        .set(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key, "true")
+        .set(RapidsConf.OPTIMIZER_ENABLED.key, "true")
+        .set(RapidsConf.TEST_ALLOWED_NONGPU.key,
+          "ProjectExec,SortMergeJoinExec,SortExec,Alias,Cast,LessThan")
+
+    var plans: ListBuffer[SparkPlanMeta[SparkPlan]] =
+      new ListBuffer[SparkPlanMeta[SparkPlan]]()
+    GpuOverrides.addListener(
+      (plan: SparkPlanMeta[SparkPlan],
+          _: SparkPlan,
+          _: Seq[Optimization]) => {
+        plans += plan
+      })
+
+    withGpuSparkSession(spark => {
+      val df1: DataFrame = createQuery(spark).alias("l")
+      val df2: DataFrame = createQuery(spark).alias("r")
+      val df = df1.join(df2,
+        col("l.more_strings_1").equalTo(col("r.more_strings_2")))
+      df.collect()
+    }, conf)
+
+    val accum = new ListBuffer[SparkPlanMeta[_]]()
+    plans.foreach(collectPlansWithRowCount(_, accum))
+
+    val summary = accum
+        .map(plan => plan.wrapped.getClass.getSimpleName -> plan.estimatedOutputRows.get)
+        .distinct
+        .sorted
+
+    // due to the concurrent nature of adaptive execution, the results are not deterministic
+    // so we just check that we do see row counts for multiple broadcast exchanges
+
+    val broadcastExchanges = summary
+        .filter(_._1 == "BroadcastExchangeExec")
+
+    assert(broadcastExchanges.nonEmpty)
+    assert(broadcastExchanges.forall(_._2.toLong > 0))
+  }
+
+  private def collectPlansWithRowCount(
+      plan: SparkPlanMeta[_],
+      accum: ListBuffer[SparkPlanMeta[_]]): Unit = {
+    if (plan.estimatedOutputRows.exists(_ > 0)) {
+      accum += plan
+    }
+    plan.childPlans.foreach(collectPlansWithRowCount(_, accum))
+  }
+
   private def createQuery(spark: SparkSession) = {
     val df1 = nullableStringsDf(spark)
         .repartition(2)
@@ -409,12 +514,4 @@ class CostBasedOptimizerSuite extends SparkQueryCompareTestSuite with BeforeAndA
     df
   }
 
-  private def addListener(optimizations: ListBuffer[Optimization]): Unit = {
-    GpuOverrides.addListener(
-      (plan: SparkPlanMeta[SparkPlan],
-          sparkPlan: SparkPlan,
-          costOptimizations: Seq[Optimization]) => {
-        optimizations.appendAll(costOptimizations)
-      })
-  }
 }

tests/src/test/scala/com/nvidia/spark/rapids/SparkQueryCompareTestSuite.scala

@@ -1811,4 +1811,24 @@ trait SparkQueryCompareTestSuite extends FunSuite with Arm {
     }
   }
 
+  /** most of the AQE tests requires Spark 3.0.1 or later */
+  def assumeSpark301orLater =
+    assume(cmpSparkVersion(3, 0, 1) >= 0)
+
+  def assumePriorToSpark320 =
+    assume(cmpSparkVersion(3, 2, 0) < 0)
+
+  def cmpSparkVersion(major: Int, minor: Int, bugfix: Int): Int = {
+    val sparkShimVersion = ShimLoader.getSparkShims.getSparkShimVersion
+    val (sparkMajor, sparkMinor, sparkBugfix) = sparkShimVersion match {
+      case SparkShimVersion(a, b, c) => (a, b, c)
+      case DatabricksShimVersion(a, b, c) => (a, b, c)
+      case EMRShimVersion(a, b, c) => (a, b, c)
+    }
+    val fullVersion = ((major.toLong * 1000) + minor) * 1000 + bugfix
+    val sparkFullVersion = ((sparkMajor.toLong * 1000) + sparkMinor) * 1000 + sparkBugfix
+    sparkFullVersion.compareTo(fullVersion)
+  }
+
+
 }