add checks to microbenchmarks for plan running natively in Comet

Author: andygrove

spark/src/test/scala/org/apache/spark/sql/benchmark/CometBenchmarkBase.scala

@@ -31,14 +31,18 @@ import org.apache.parquet.crypto.keytools.mocks.InMemoryKMS
 import org.apache.spark.SparkConf
 import org.apache.spark.benchmark.Benchmark
 import org.apache.spark.sql.{DataFrame, DataFrameWriter, Row, SparkSession}
+import org.apache.spark.sql.comet._
+import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec
+import org.apache.spark.sql.execution.{ColumnarToRowExec, InputAdapter, SparkPlan, WholeStageCodegenExec}
+import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
 import org.apache.spark.sql.execution.benchmark.SqlBasedBenchmark
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types.DecimalType
 
 import org.apache.comet.CometConf
 import org.apache.comet.CometSparkSessionExtensions
 
-trait CometBenchmarkBase extends SqlBasedBenchmark {
+trait CometBenchmarkBase extends SqlBasedBenchmark with AdaptiveSparkPlanHelper {
   override def getSparkSession: SparkSession = {
     val conf = new SparkConf()
       .setAppName("CometReadBenchmark")
@@ -155,9 +159,53 @@ trait CometBenchmarkBase extends SqlBasedBenchmark {
       }
     }
 
+    // Check that the plan is fully Comet native before running the benchmark
+    withSQLConf(cometExecConfigs.toSeq: _*) {
+      val df = spark.sql(query)
+      val plan = stripAQEPlan(df.queryExecution.executedPlan)
+      findFirstNonCometOperator(plan) match {
+        case Some(op) =>
+          // scalastyle:off println
+          println()
+          println("=" * 80)
+          println("WARNING: Benchmark plan is NOT fully Comet native!")
+          println(s"First non-Comet operator: ${op.nodeName}")
+          println("=" * 80)
+          println("Query plan:")
+          println(plan.treeString)
+          println("=" * 80)
+          println()
+        // scalastyle:on println
+        case None =>
+        // All operators are Comet native, no warning needed
+      }
+    }
+
     benchmark.run()
   }
 
+  /**
+   * Finds the first non-Comet operator in the plan, if any. This is used to verify that
+   * benchmarks are running fully on Comet native when expected.
+   *
+   * Based on CometTestBase.findFirstNonCometOperator.
+   */
+  protected def findFirstNonCometOperator(plan: SparkPlan): Option[SparkPlan] = {
+    plan.foreach {
+      case _: CometNativeScanExec | _: CometScanExec | _: CometBatchScanExec |
+          _: CometIcebergNativeScanExec =>
+      case _: CometSinkPlaceHolder | _: CometScanWrapper =>
+      case _: CometColumnarToRowExec =>
+      case _: CometSparkToColumnarExec =>
+      case _: CometExec | _: CometShuffleExchangeExec =>
+      case _: CometBroadcastExchangeExec =>
+      case _: WholeStageCodegenExec | _: ColumnarToRowExec | _: InputAdapter =>
+      case op =>
+        return Some(op)
+    }
+    None
+  }
+
   protected def prepareTable(dir: File, df: DataFrame, partition: Option[String] = None): Unit = {
     val testDf = if (partition.isDefined) {
       df.write.partitionBy(partition.get)