Fix unit tests when AQE is enabled

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

@@ -16,6 +16,8 @@
 
 package com.nvidia.spark.rapids
 
+import com.nvidia.spark.rapids.TestUtils.{findOperator, operatorCount}
+
 import org.apache.spark.SparkConf
 import org.apache.spark.sql.execution.joins.HashJoin
 import org.apache.spark.sql.functions.broadcast
@@ -35,13 +37,14 @@ class BroadcastHashJoinSuite extends SparkQueryCompareTestSuite {
       val df5 = df4.join(df3, Seq("longs"), "inner")
 
       val plan = df5.queryExecution.executedPlan
+      // execute the plan so that the final adaptive plan is available when AQE is on
+      df5.collect()
+
+      val bhjCount = operatorCount(plan, ShimLoader.getSparkShims.isGpuBroadcastHashJoin)
+      assert(bhjCount.size === 1)
 
-      assert(plan.collect {
-        case p if ShimLoader.getSparkShims.isGpuBroadcastHashJoin(p) => p
-      }.size === 1)
-      assert(plan.collect {
-        case p if ShimLoader.getSparkShims.isGpuShuffledHashJoin(p) => p
-      }.size === 1)
+      val shjCount = operatorCount(plan, ShimLoader.getSparkShims.isGpuShuffledHashJoin)
+      assert(shjCount.size === 1)
     }, conf)
   }
 
@@ -52,17 +55,21 @@ class BroadcastHashJoinSuite extends SparkQueryCompareTestSuite {
 
       for (name <- Seq("BROADCAST", "BROADCASTJOIN", "MAPJOIN")) {
         val plan1 = spark.sql(s"SELECT /*+ $name(t) */ * FROM t JOIN u ON t.longs = u.longs")
-          .queryExecution.executedPlan
         val plan2 = spark.sql(s"SELECT /*+ $name(u) */ * FROM t JOIN u ON t.longs = u.longs")
-          .queryExecution.executedPlan
 
-        val res1 = plan1.find(ShimLoader.getSparkShims.isGpuBroadcastHashJoin(_))
-        val res2 = plan2.find(ShimLoader.getSparkShims.isGpuBroadcastHashJoin(_))
+        val initialPlan1 = plan1.queryExecution.executedPlan
+        // execute the plan so that the final adaptive plan is available when AQE is on
+        plan1.collect()
+        val finalPlan1 = findOperator(initialPlan1, ShimLoader.getSparkShims.isGpuBroadcastHashJoin)
+        assert(ShimLoader.getSparkShims.getBuildSide
+        (finalPlan1.get.asInstanceOf[HashJoin]).toString == "GpuBuildLeft")
 
-        assert(ShimLoader.getSparkShims.getBuildSide(res1.get.asInstanceOf[HashJoin]).toString ==
-          "GpuBuildLeft")
-        assert(ShimLoader.getSparkShims.getBuildSide(res2.get.asInstanceOf[HashJoin]).toString ==
-          "GpuBuildRight")
+        val initialPlan2 = plan2.queryExecution.executedPlan
+        // execute the plan so that the final adaptive plan is available when AQE is on
+        plan2.collect()
+        val finalPlan2 = findOperator(initialPlan2, ShimLoader.getSparkShims.isGpuBroadcastHashJoin)
+        assert(ShimLoader.getSparkShims.
+          getBuildSide(finalPlan2.get.asInstanceOf[HashJoin]).toString == "GpuBuildRight")
       }
     })
   }

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

@@ -16,6 +16,7 @@
 
 package com.nvidia.spark.rapids
 
+import com.nvidia.spark.rapids.TestUtils.{findOperator, getFinalPlan}
 import org.scalatest.FunSuite
 
 import org.apache.spark.SparkConf
@@ -69,7 +70,10 @@ class HashSortOptimizeSuite extends FunSuite {
       val df2 = buildDataFrame2(spark)
       val rdf = df1.join(df2, df1("a") === df2("x"))
       val plan = rdf.queryExecution.executedPlan
-      val joinNode = plan.find(ShimLoader.getSparkShims.isGpuBroadcastHashJoin(_))
+      // execute the plan so that the final adaptive plan is available when AQE is on
+      rdf.collect()
+
+      val joinNode = findOperator(plan, ShimLoader.getSparkShims.isGpuBroadcastHashJoin(_))
       assert(joinNode.isDefined, "No broadcast join node found")
       validateOptimizeSort(plan, joinNode.get)
     })
@@ -82,7 +86,9 @@ class HashSortOptimizeSuite extends FunSuite {
       val df2 = buildDataFrame2(spark)
       val rdf = df1.join(df2, df1("a") === df2("x"))
       val plan = rdf.queryExecution.executedPlan
-      val joinNode = plan.find(ShimLoader.getSparkShims.isGpuShuffledHashJoin(_))
+      // execute the plan so that the final adaptive plan is available when AQE is on
+      rdf.collect()
+      val joinNode = findOperator(plan, ShimLoader.getSparkShims.isGpuShuffledHashJoin(_))
       assert(joinNode.isDefined, "No broadcast join node found")
       validateOptimizeSort(plan, joinNode.get)
     })
@@ -106,7 +112,10 @@ class HashSortOptimizeSuite extends FunSuite {
       val df2 = buildDataFrame2(spark)
       val rdf = df1.join(df2, df1("a") === df2("x")).orderBy(df1("a"))
       val plan = rdf.queryExecution.executedPlan
-      val numSorts = plan.map {
+      // Get the final executed plan when AQE is either enabled or disabled.
+      val finalPlan = getFinalPlan(plan)
+
+      val numSorts = finalPlan.map {
         case _: SortExec | _: GpuSortExec => 1
         case _ => 0
       }.sum

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

@@ -20,9 +20,12 @@ import java.io.File
 
 import ai.rapids.cudf.{ColumnVector, DType, Table}
 import org.scalatest.Assertions
+import scala.collection.mutable.ListBuffer
 
 import org.apache.spark.SparkConf
 import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, BroadcastQueryStageExec, ShuffleQueryStageExec}
 import org.apache.spark.sql.rapids.GpuShuffleEnv
 import org.apache.spark.sql.vectorized.ColumnarBatch
 
@@ -51,6 +54,47 @@ object TestUtils extends Assertions with Arm {
     }
   }
 
+  /** Recursively check if the predicate matches in the given plan */
+  def findOperator(plan: SparkPlan, predicate: SparkPlan => Boolean): Option[SparkPlan] = {
+    plan match {
+      case _ if predicate(plan) => Some(plan)
+      case a: AdaptiveSparkPlanExec => findOperator(a.executedPlan, predicate)
+      case qs: BroadcastQueryStageExec => findOperator(qs.broadcast, predicate)
+      case qs: ShuffleQueryStageExec => findOperator(qs.shuffle, predicate)
+      case other => other.children.flatMap(p => findOperator(p, predicate)).headOption
+    }
+  }
+
+  /** Return list of  matching predicates present in the plan */
+  def operatorCount(plan: SparkPlan, predicate: SparkPlan => Boolean): Seq[SparkPlan] = {
+    def recurse(
+      plan: SparkPlan,
+      predicate: SparkPlan => Boolean,
+      accum: ListBuffer[SparkPlan]): Seq[SparkPlan] = {
+      plan match {
+        case _ if predicate(plan) =>
+          accum += plan
+          plan.children.flatMap(p => recurse(p, predicate, accum)).headOption
+        case a: AdaptiveSparkPlanExec => recurse(a.executedPlan, predicate, accum)
+        case qs: BroadcastQueryStageExec => recurse(qs.broadcast, predicate, accum)
+        case qs: ShuffleQueryStageExec => recurse(qs.shuffle, predicate, accum)
+        case other => other.children.flatMap(p => recurse(p, predicate, accum)).headOption
+      }
+      accum
+    }
+
+    recurse(plan, predicate, new ListBuffer[SparkPlan]())
+  }
+
+  /** Return final executed plan */
+  def getFinalPlan(plan: SparkPlan): SparkPlan = {
+    plan match {
+      case a: AdaptiveSparkPlanExec =>
+        a.executedPlan
+      case _ => plan
+    }
+  }
+
   /** Compre the equality of two `ColumnVector` instances */
   def compareColumns(expected: ColumnVector, actual: ColumnVector): Unit = {
     assertResult(expected.getType)(actual.getType)