test

Author: andygrove

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

@@ -21,6 +21,7 @@ package org.apache.comet.cost
 
 import scala.jdk.CollectionConverters._
 
+import org.apache.spark.internal.Logging
 import org.apache.spark.sql.comet.{CometColumnarToRowExec, CometPlan, CometProjectExec}
 import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometNativeShuffle, CometShuffleExchangeExec}
 import org.apache.spark.sql.execution.SparkPlan
@@ -36,23 +37,24 @@ trait CometCostModel {
   def estimateCost(plan: SparkPlan): CometCostEstimate
 }
 
-class DefaultCometCostModel extends CometCostModel {
+class DefaultCometCostModel extends CometCostModel with Logging {
 
   // optimistic default of 2x acceleration
   private val defaultAcceleration = 2.0
 
   override def estimateCost(plan: SparkPlan): CometCostEstimate = {
+
+    logTrace(s"estimateCost for $plan")
+
     // Walk the entire plan tree and accumulate costs
     var totalAcceleration = 0.0
     var operatorCount = 0
 
     def collectOperatorCosts(node: SparkPlan): Unit = {
       val operatorCost = estimateOperatorCost(node)
-      // scalastyle:off println
-      println(
-        s"[CostModel] Operator: ${node.getClass.getSimpleName}, " +
+      logTrace(
+        s"Operator: ${node.getClass.getSimpleName}, " +
           s"Cost: ${operatorCost.acceleration}")
-      // scalastyle:on println
       totalAcceleration += operatorCost.acceleration
       operatorCount += 1
 
@@ -70,11 +72,9 @@ class DefaultCometCostModel extends CometCostModel {
       1.0 // No acceleration if no operators
     }
 
-    // scalastyle:off println
-    println(
-      s"[CostModel] Plan: ${plan.getClass.getSimpleName}, Total operators: $operatorCount, " +
+    logTrace(
+      s"Plan: ${plan.getClass.getSimpleName}, Total operators: $operatorCount, " +
         s"Average acceleration: $averageAcceleration")
-    // scalastyle:on println
 
     CometCostEstimate(averageAcceleration)
   }
@@ -83,29 +83,21 @@ class DefaultCometCostModel extends CometCostModel {
   private def estimateOperatorCost(plan: SparkPlan): CometCostEstimate = {
     val result = plan match {
       case op: CometProjectExec =>
-        // scalastyle:off println
-        println(s"[CostModel] CometProjectExec found - evaluating expressions")
-        // scalastyle:on println
+        logTrace("CometProjectExec found - evaluating expressions")
         // Cast nativeOp to Operator and extract projection expressions
         val operator = op.nativeOp.asInstanceOf[OperatorOuterClass.Operator]
         val projection = operator.getProjection
         val expressions = projection.getProjectListList.asScala
-        // scalastyle:off println
-        println(s"[CostModel] Found ${expressions.length} expressions in projection")
-        // scalastyle:on println
+        logTrace(s"Found ${expressions.length} expressions in projection")
 
         val costs = expressions.map { expr =>
           val cost = estimateCometExpressionCost(expr)
-          // scalastyle:off println
-          println(s"[CostModel] Expression cost: $cost")
-          // scalastyle:on println
+          logTrace(s"Expression cost: $cost")
           cost
         }
         val total = costs.sum
         val average = total / expressions.length.toDouble
-        // scalastyle:off println
-        println(s"[CostModel] CometProjectExec total cost: $total, average: $average")
-        // scalastyle:on println
+        logTrace(s"CometProjectExec total cost: $total, average: $average")
         CometCostEstimate(average)
 
       case op: CometShuffleExchangeExec =>
@@ -121,38 +113,28 @@ class DefaultCometCostModel extends CometCostModel {
       case _: CometColumnarToRowExec =>
         CometCostEstimate(1.0)
       case _: CometPlan =>
-        // scalastyle:off println
-        println(s"[CostModel] Generic CometPlan: ${plan.getClass.getSimpleName}")
-        // scalastyle:on println
+        logTrace(s"Generic CometPlan: ${plan.getClass.getSimpleName}")
         CometCostEstimate(defaultAcceleration)
       case _ =>
-        // scalastyle:off println
-        println(s"[CostModel] Non-Comet operator: ${plan.getClass.getSimpleName}")
-        // scalastyle:on println
+        logTrace(s"Non-Comet operator: ${plan.getClass.getSimpleName}")
         // Spark operator
         CometCostEstimate(1.0)
     }
 
-    // scalastyle:off println
-    println(s"[CostModel] ${plan.getClass.getSimpleName} -> acceleration: ${result.acceleration}")
-    // scalastyle:on println
+    logTrace(s"${plan.getClass.getSimpleName} -> acceleration: ${result.acceleration}")
     result
   }
 
   /** Estimate the cost of a Comet protobuf expression */
   private def estimateCometExpressionCost(expr: ExprOuterClass.Expr): Double = {
     val result = expr.getExprStructCase match {
       // Handle specialized expression types
-      case ExprOuterClass.Expr.ExprStructCase.SUBSTRING =>
-        // scalastyle:off println
-        println(s"[CostModel] Expression: SUBSTRING -> 6.3")
-        // scalastyle:on println
-        6.3
+      case ExprOuterClass.Expr.ExprStructCase.SUBSTRING => 6.3
 
       // Handle generic scalar functions
       case ExprOuterClass.Expr.ExprStructCase.SCALARFUNC =>
         val funcName = expr.getScalarFunc.getFunc
-        val cost = funcName match {
+        funcName match {
           // String expression numbers from CometStringExpressionBenchmark
           case "ascii" => 0.6
           case "octet_length" => 0.6
@@ -171,17 +153,11 @@ class DefaultCometCostModel extends CometCostModel {
           case "translate" => 0.8
           case _ => defaultAcceleration
         }
-        // scalastyle:off println
-        println(s"[CostModel] Expression: SCALARFUNC($funcName) -> $cost")
-        // scalastyle:on println
-        cost
 
       case _ =>
-        // scalastyle:off println
-        println(
-          s"[CostModel] Expression: Unknown type ${expr.getExprStructCase} -> " +
+        logTrace(
+          s"Expression: Unknown type ${expr.getExprStructCase} -> " +
             s"$defaultAcceleration")
-        // scalastyle:on println
         defaultAcceleration
     }
     result

spark/src/test/scala/org/apache/comet/CometCostModelSuite.scala

@@ -19,7 +19,7 @@
 
 package org.apache.comet
 
-import org.apache.spark.sql.{CometTestBase, DataFrame}
+import org.apache.spark.sql.CometTestBase
 import org.apache.spark.sql.comet.CometProjectExec
 import org.apache.spark.sql.execution.ProjectExec
 import org.apache.spark.sql.execution.SparkPlan
@@ -79,7 +79,7 @@ class CometCostModelSuite extends CometTestBase {
       "spark.sql.adaptive.advisoryPartitionSizeInBytes" -> "1KB",
       "spark.sql.adaptive.coalescePartitions.minPartitionSize" -> "1B") {
 
-      println(s"\n=== CBO Configuration ===")
+      println("\n=== CBO Configuration ===")
       println(s"COMET_ENABLED: ${spark.conf.get(CometConf.COMET_ENABLED.key)}")
       println(s"COMET_COST_BASED_OPTIMIZATION_ENABLED: ${spark.conf.get(
           CometConf.COMET_COST_BASED_OPTIMIZATION_ENABLED.key)}")
@@ -148,13 +148,13 @@ class CometCostModelSuite extends CometTestBase {
       expectedClass: Class[_],
       message: String): Unit = {
 
-    println(s"\n=== Executing Query ===")
+    println("\n=== Executing Query ===")
     println(s"Query: $query")
     println(s"Expected class: ${expectedClass.getSimpleName}")
 
     val result = sql(query)
 
-    println(s"\n=== Pre-execution Plans ===")
+    println("\n=== Pre-execution Plans ===")
     println("Logical Plan:")
     println(result.queryExecution.logical)
     println("\nOptimized Plan:")
@@ -164,7 +164,7 @@ class CometCostModelSuite extends CometTestBase {
 
     result.collect() // Materialize the plan
 
-    println(s"\n=== Post-execution Plans ===")
+    println("\n=== Post-execution Plans ===")
     println("Executed Plan (with AQE wrappers):")
     println(result.queryExecution.executedPlan)
 
@@ -178,7 +178,7 @@ class CometCostModelSuite extends CometTestBase {
 
     val hasProjectExec = findProjectExec(executedPlan)
 
-    println(s"\n=== Project Analysis ===")
+    println("\n=== Project Analysis ===")
     println(s"Found project exec: ${hasProjectExec.isDefined}")
     if (hasProjectExec.isDefined) {
       println(s"Actual class: ${hasProjectExec.get.getClass.getSimpleName}")
@@ -191,7 +191,7 @@ class CometCostModelSuite extends CometTestBase {
       expectedClass.isInstance(hasProjectExec.get),
       s"$message, got ${hasProjectExec.get.getClass.getSimpleName}")
 
-    println(s"=== Test PASSED ===\n")
+    println("=== Test PASSED ===\n")
   }
 
   /** Helper method to find ProjectExec or CometProjectExec in the plan tree */