refactor

Author: andygrove

spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala

@@ -38,6 +38,7 @@ import org.apache.parquet.schema.{MessageType, MessageTypeParser}
 import org.apache.spark._
 import org.apache.spark.internal.config.{MEMORY_OFFHEAP_ENABLED, MEMORY_OFFHEAP_SIZE, SHUFFLE_MANAGER}
 import org.apache.spark.sql.comet._
+import org.apache.spark.sql.comet.CometPlanChecker
 import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometNativeShuffle, CometShuffleExchangeExec}
 import org.apache.spark.sql.execution._
 import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
@@ -58,7 +59,8 @@ abstract class CometTestBase
     with BeforeAndAfterEach
     with AdaptiveSparkPlanHelper
     with ShimCometSparkSessionExtensions
-    with ShimCometTestBase {
+    with ShimCometTestBase
+    with CometPlanChecker {
   import testImplicits._
 
   protected val shuffleManager: String =
@@ -396,26 +398,6 @@ abstract class CometTestBase
     checkPlanNotMissingInput(plan)
   }
 
-  protected def findFirstNonCometOperator(
-      plan: SparkPlan,
-      excludedClasses: Class[_]*): Option[SparkPlan] = {
-    val wrapped = wrapCometSparkToColumnar(plan)
-    wrapped.foreach {
-      case _: CometNativeScanExec | _: CometScanExec | _: CometBatchScanExec |
-          _: CometIcebergNativeScanExec =>
-      case _: CometSinkPlaceHolder | _: CometScanWrapper =>
-      case _: CometColumnarToRowExec =>
-      case _: CometSparkToColumnarExec =>
-      case _: CometExec | _: CometShuffleExchangeExec =>
-      case _: CometBroadcastExchangeExec =>
-      case _: WholeStageCodegenExec | _: ColumnarToRowExec | _: InputAdapter =>
-      case op if !excludedClasses.exists(c => c.isAssignableFrom(op.getClass)) =>
-        return Some(op)
-      case _ =>
-    }
-    None
-  }
-
   // checks the plan node has no missing inputs
   // such nodes represented in plan with exclamation mark !
   // example: !CometWindowExec
@@ -449,14 +431,6 @@ abstract class CometTestBase
     }
   }
 
-  /** Wraps the CometRowToColumn as ScanWrapper, so the child operators will not be checked */
-  private def wrapCometSparkToColumnar(plan: SparkPlan): SparkPlan = {
-    plan.transformDown {
-      // don't care the native operators
-      case p: CometSparkToColumnarExec => CometScanWrapper(null, p)
-    }
-  }
-
   private var _spark: SparkSessionType = _
   override protected implicit def spark: SparkSessionType = _spark
   protected implicit def sqlContext: SQLContext = _spark.sqlContext

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

@@ -31,9 +31,7 @@ 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.comet.CometPlanChecker
 import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
 import org.apache.spark.sql.execution.benchmark.SqlBasedBenchmark
 import org.apache.spark.sql.internal.SQLConf
@@ -42,7 +40,10 @@ import org.apache.spark.sql.types.DecimalType
 import org.apache.comet.CometConf
 import org.apache.comet.CometSparkSessionExtensions
 
-trait CometBenchmarkBase extends SqlBasedBenchmark with AdaptiveSparkPlanHelper {
+trait CometBenchmarkBase
+    extends SqlBasedBenchmark
+    with AdaptiveSparkPlanHelper
+    with CometPlanChecker {
   override def getSparkSession: SparkSession = {
     val conf = new SparkConf()
       .setAppName("CometReadBenchmark")
@@ -163,28 +164,6 @@ trait CometBenchmarkBase extends SqlBasedBenchmark with AdaptiveSparkPlanHelper
     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)

spark/src/test/scala/org/apache/spark/sql/comet/CometPlanChecker.scala

@@ -0,0 +1,68 @@
+/*
+ * 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.comet
+
+import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec
+import org.apache.spark.sql.execution.{ColumnarToRowExec, InputAdapter, SparkPlan, WholeStageCodegenExec}
+
+/**
+ * Trait providing utilities to check if a Spark plan is fully running on Comet native operators.
+ * Used by both CometTestBase and CometBenchmarkBase.
+ */
+trait CometPlanChecker {
+
+  /**
+   * Finds the first non-Comet operator in the plan, if any.
+   *
+   * @param plan
+   *   The SparkPlan to check
+   * @param excludedClasses
+   *   Classes to exclude from the check (these are allowed to be non-Comet)
+   * @return
+   *   Some(operator) if a non-Comet operator is found, None otherwise
+   */
+  protected def findFirstNonCometOperator(
+      plan: SparkPlan,
+      excludedClasses: Class[_]*): Option[SparkPlan] = {
+    val wrapped = wrapCometSparkToColumnar(plan)
+    wrapped.foreach {
+      case _: CometNativeScanExec | _: CometScanExec | _: CometBatchScanExec |
+          _: CometIcebergNativeScanExec =>
+      case _: CometSinkPlaceHolder | _: CometScanWrapper =>
+      case _: CometColumnarToRowExec =>
+      case _: CometSparkToColumnarExec =>
+      case _: CometExec | _: CometShuffleExchangeExec =>
+      case _: CometBroadcastExchangeExec =>
+      case _: WholeStageCodegenExec | _: ColumnarToRowExec | _: InputAdapter =>
+      case op if !excludedClasses.exists(c => c.isAssignableFrom(op.getClass)) =>
+        return Some(op)
+      case _ =>
+    }
+    None
+  }
+
+  /** Wraps the CometSparkToColumnar as ScanWrapper, so the child operators will not be checked */
+  private def wrapCometSparkToColumnar(plan: SparkPlan): SparkPlan = {
+    plan.transformDown {
+      // don't care the native operators
+      case p: CometSparkToColumnarExec => CometScanWrapper(null, p)
+    }
+  }
+}