remove AQE integration

Author: andygrove

spark/src/main/scala/org/apache/comet/cost/CometCostEvaluator.scala

@@ -49,6 +49,7 @@ import org.apache.spark.sql.types._
 import org.apache.comet.{CometConf, CometExplainInfo, ExtendedExplainInfo}
 import org.apache.comet.CometConf.{COMET_SPARK_TO_ARROW_ENABLED, COMET_SPARK_TO_ARROW_SUPPORTED_OPERATOR_LIST}
 import org.apache.comet.CometSparkSessionExtensions._
+import org.apache.comet.cost.DefaultCometCostModel
 import org.apache.comet.rules.CometExecRule.allExecs
 import org.apache.comet.serde.{CometOperatorSerde, Compatible, Incompatible, OperatorOuterClass, Unsupported}
 import org.apache.comet.serde.operator._
@@ -344,7 +345,19 @@ case class CometExecRule(session: SparkSession) extends Rule[SparkPlan] {
   }
 
   override def apply(plan: SparkPlan): SparkPlan = {
-    val newPlan = _apply(plan)
+    val candidatePlan = _apply(plan)
+
+    // TODO load cost model via config and reflection
+    val costModel = new DefaultCometCostModel
+    val costBefore = costModel.estimateCost(plan)
+    val costAfter = costModel.estimateCost(candidatePlan)
+
+    val newPlan = if (costAfter.acceleration > costBefore.acceleration) {
+      candidatePlan
+    } else {
+      plan
+    }
+
     if (showTransformations && !newPlan.fastEquals(plan)) {
       logInfo(s"""
            |=== Applying Rule $ruleName ===

spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala

@@ -28,7 +28,7 @@ import org.apache.spark.internal.Logging
 import org.apache.spark.internal.config.{EXECUTOR_MEMORY, EXECUTOR_MEMORY_OVERHEAD, EXECUTOR_MEMORY_OVERHEAD_FACTOR}
 import org.apache.spark.sql.internal.StaticSQLConf
 
-import org.apache.comet.CometConf.{COMET_COST_BASED_OPTIMIZATION_ENABLED, COMET_ONHEAP_ENABLED}
+import org.apache.comet.CometConf.COMET_ONHEAP_ENABLED
 import org.apache.comet.CometSparkSessionExtensions
 
 /**
@@ -57,15 +57,6 @@ class CometDriverPlugin extends DriverPlugin with Logging with ShimCometDriverPl
     // register CometSparkSessionExtensions if it isn't already registered
     CometDriverPlugin.registerCometSessionExtension(sc.conf)
 
-    // Enable cost-based optimization if configured
-    if (sc.getConf.getBoolean(COMET_COST_BASED_OPTIMIZATION_ENABLED.key, false)) {
-      // Set the custom cost evaluator for Spark's adaptive query execution
-      sc.conf.set(
-        "spark.sql.adaptive.customCostEvaluatorClass",
-        "org.apache.comet.cost.CometCostEvaluator")
-      logInfo("Enabled Comet cost-based optimization with CometCostEvaluator")
-    }
-
     if (CometSparkSessionExtensions.shouldOverrideMemoryConf(sc.getConf)) {
       val execMemOverhead = if (sc.getConf.contains(EXECUTOR_MEMORY_OVERHEAD.key)) {
         sc.getConf.getSizeAsMb(EXECUTOR_MEMORY_OVERHEAD.key)

spark/src/main/scala/org/apache/spark/Plugins.scala

@@ -73,8 +73,6 @@ class CometCostModelSuite extends CometTestBase {
       CometConf.COMET_EXEC_PROJECT_ENABLED.key -> "true",
       CometConf.COMET_EXEC_AGGREGATE_ENABLED.key -> "true", // Enable aggregation for GROUP BY
       CometConf.COMET_EXEC_HASH_JOIN_ENABLED.key -> "true", // Enable joins
-      // Manually set the custom cost evaluator since plugin might not be loaded
-      "spark.sql.adaptive.customCostEvaluatorClass" -> "org.apache.comet.cost.CometCostEvaluator",
       // Lower AQE thresholds to ensure it triggers on small test data
       "spark.sql.adaptive.advisoryPartitionSizeInBytes" -> "1KB",
       "spark.sql.adaptive.coalescePartitions.minPartitionSize" -> "1B") {