Convert prints to logs.

Author: mbutrovich

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

@@ -301,100 +301,60 @@ case class CometExecRule(session: SparkSession)
         }
     }
 
-    // scalastyle:off println
     plan.transformUp { case op =>
-      val hasSubqueryExpr = op.expressions.exists(_.exists {
-        case _: InSubqueryExec => true
-        case _ => false
-      })
-      if (hasSubqueryExpr) {
-        println(s"[RULE-DEBUG] convertNode on ${op.getClass.getSimpleName} " +
-          s"which HAS InSubqueryExec expressions")
-        op.expressions.foreach { expr =>
-          expr.foreach {
-            case sub: InSubqueryExec =>
-              println(s"[RULE-DEBUG]   InSubqueryExec.plan: ${sub.plan.getClass.getSimpleName}")
-              sub.plan match {
-                case sb: SubqueryBroadcastExec =>
-                  println(s"[RULE-DEBUG]   SubqueryBroadcast.child: " +
-                    s"${sb.child.getClass.getSimpleName}")
-                  sb.child match {
-                    case b: BroadcastExchangeExec =>
-                      println(s"[RULE-DEBUG]   BroadcastExchange.child: " +
-                        s"${b.child.getClass.getSimpleName}")
-                      println(s"[RULE-DEBUG]   BroadcastExchange.child is CometNative? " +
-                        s"${b.child.isInstanceOf[CometNativeExec]}")
-                      println(s"[RULE-DEBUG]   BroadcastExchange.children all CometNative? " +
-                        s"${b.children.forall(_.isInstanceOf[CometNativeExec])}")
-                    case other =>
-                      println(s"[RULE-DEBUG]   SubqueryBroadcast.child is: " +
-                        s"${other.getClass.getSimpleName}")
-                  }
-                case other =>
-                  println(s"[RULE-DEBUG]   sub.plan is: ${other.getClass.getSimpleName}")
-              }
-            case _ =>
-          }
-        }
-      }
       val converted = convertNode(op)
       // Replace SubqueryBroadcastExec with CometSubqueryBroadcastExec in DPP expressions
       // when the broadcast child has a Comet plan underneath. This enables exchange reuse
       // between the DPP subquery and the join's CometBroadcastExchangeExec because both
       // will have the same CometBroadcastExchangeExec type and canonical form.
       convertSubqueryBroadcasts(converted)
     }
-    // scalastyle:on println
   }
 
   /**
    * Replace SubqueryBroadcastExec with CometSubqueryBroadcastExec in a node's expressions.
    *
-   * When CometExecRule converts BroadcastExchangeExec to CometBroadcastExchangeExec on the
-   * join side, the DPP subquery still references the original BroadcastExchangeExec.
-   * ReuseExchangeAndSubquery (which runs after Comet rules) can't match them because they
-   * have different types. By replacing SubqueryBroadcastExec with CometSubqueryBroadcastExec
-   * (which wraps a CometBroadcastExchangeExec), both sides have the same exchange type and
-   * reuse works.
+   * When CometExecRule converts BroadcastExchangeExec to CometBroadcastExchangeExec on the join
+   * side, the DPP subquery still references the original BroadcastExchangeExec.
+   * ReuseExchangeAndSubquery (which runs after Comet rules) can't match them because they have
+   * different types. By replacing SubqueryBroadcastExec with CometSubqueryBroadcastExec (which
+   * wraps a CometBroadcastExchangeExec), both sides have the same exchange type and reuse works.
    *
-   * The BroadcastExchangeExec in the subquery has a CometNativeColumnarToRowExec child
-   * (inserted by ApplyColumnarRulesAndInsertTransitions because BroadcastExchangeExec expects
-   * row input). We strip this transition and create CometBroadcastExchangeExec with the
-   * underlying Comet plan directly.
+   * The BroadcastExchangeExec in the subquery has a CometNativeColumnarToRowExec child (inserted
+   * by ApplyColumnarRulesAndInsertTransitions because BroadcastExchangeExec expects row input).
+   * We strip this transition and create CometBroadcastExchangeExec with the underlying Comet plan
+   * directly.
    */
   private def convertSubqueryBroadcasts(plan: SparkPlan): SparkPlan = {
-    plan.transformExpressionsUp {
-      case inSub: InSubqueryExec =>
-        inSub.plan match {
-          case sub: SubqueryBroadcastExec =>
-            sub.child match {
-              case b: BroadcastExchangeExec =>
-                // The BroadcastExchangeExec child is CometNativeColumnarToRowExec wrapping
-                // a Comet plan. Strip the row transition to get the columnar Comet plan.
-                val cometChild = b.child match {
-                  case c2r: CometNativeColumnarToRowExec => c2r.child
-                  case other => other
-                }
-                if (cometChild.isInstanceOf[CometNativeExec]) {
-                  // scalastyle:off println
-                  println(s"[RULE-DEBUG] Converting SubqueryBroadcastExec to " +
-                    s"CometSubqueryBroadcastExec, cometChild=${cometChild.getClass.getSimpleName}")
-                  // scalastyle:on println
-                  val cometBroadcast = CometBroadcastExchangeExec(
-                    b, b.output, b.mode, cometChild)
-                  val cometSub = CometSubqueryBroadcastExec(
-                    sub.name,
-                    getSubqueryBroadcastExecIndices(sub),
-                    sub.buildKeys,
-                    cometBroadcast)
-                  inSub.withNewPlan(cometSub)
-                } else {
-                  inSub
-                }
-              case _ => inSub
-            }
-          case _ => inSub
-        }
+    plan.transformExpressionsUp { case inSub: InSubqueryExec =>
+      inSub.plan match {
+        case sub: SubqueryBroadcastExec =>
+          sub.child match {
+            case b: BroadcastExchangeExec =>
+              // The BroadcastExchangeExec child is CometNativeColumnarToRowExec wrapping
+              // a Comet plan. Strip the row transition to get the columnar Comet plan.
+              val cometChild = b.child match {
+                case c2r: CometNativeColumnarToRowExec => c2r.child
+                case other => other
+              }
+              if (cometChild.isInstanceOf[CometNativeExec]) {
+                logInfo(
+                  s"Converting SubqueryBroadcastExec to " +
+                    s"CometSubqueryBroadcastExec for DPP exchange reuse")
+                val cometBroadcast = CometBroadcastExchangeExec(b, b.output, b.mode, cometChild)
+                val cometSub = CometSubqueryBroadcastExec(
+                  sub.name,
+                  getSubqueryBroadcastExecIndices(sub),
+                  sub.buildKeys,
+                  cometBroadcast)
+                inSub.withNewPlan(cometSub)
+              } else {
+                inSub
+              }
+            case _ => inSub
+          }
+        case _ => inSub
+      }
     }
   }
 

spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala

@@ -176,37 +176,20 @@ case class CometNativeScanExec(
    * partition's files (lazily, as tasks are scheduled).
    */
   @transient private lazy val serializedPartitionData: (Array[Byte], Array[Array[Byte]]) = {
-    // scalastyle:off println
-    println(s"[DPP-DEBUG] serializedPartitionData: checking partitionFilters")
-    partitionFilters.foreach {
-      case DynamicPruningExpression(e: InSubqueryExec) =>
-        println(s"[DPP-DEBUG]   InSubqueryExec plan=${e.plan.getClass.getSimpleName} " +
-          s"values empty=${e.values().isEmpty}")
-      case other =>
-        println(s"[DPP-DEBUG]   filter: ${other.getClass.getSimpleName}")
-    }
-    // scalastyle:on println
     // Ensure DPP subqueries are resolved before accessing file partitions.
     // serializedPartitionData can be triggered from findAllPlanData (via commonData) on a
     // different execution path than the standard prepare() -> executeSubqueries() flow
     // (e.g., from a BroadcastExchangeExec thread). We must resolve DPP here explicitly.
     partitionFilters.foreach {
       case DynamicPruningExpression(e: InSubqueryExec) if e.values().isEmpty =>
-        // scalastyle:off println
-        println(s"[DPP-DEBUG] calling updateResult on InSubqueryExec " +
-          s"plan=${e.plan.getClass.getSimpleName} id=${System.identityHashCode(e)}")
-        // scalastyle:on println
+        logDebug(s"Resolving DPP subquery: plan=${e.plan.getClass.getSimpleName}")
         try {
           e.updateResult()
-          // scalastyle:off println
-          println(s"[DPP-DEBUG] updateResult succeeded, values empty=${e.values().isEmpty}")
-          // scalastyle:on println
+          logDebug(s"DPP subquery resolved successfully")
         } catch {
-          // scalastyle:off println
           case ex: Exception =>
-            println(s"[DPP-DEBUG] updateResult FAILED: ${ex.getMessage}")
+            logError(s"DPP subquery resolution failed: ${ex.getMessage}")
             throw ex
-          // scalastyle:on println
         }
       case _ =>
     }
@@ -216,10 +199,9 @@ case class CometNativeScanExec(
     if (scan != null) {
       scan.partitionFilters.foreach {
         case DynamicPruningExpression(e: InSubqueryExec) if e.values().isEmpty =>
-          // scalastyle:off println
-          println(s"[DPP-DEBUG] also resolving scan's InSubqueryExec " +
-            s"plan=${e.plan.getClass.getSimpleName} id=${System.identityHashCode(e)}")
-          // scalastyle:on println
+          logDebug(
+            s"Resolving CometScanExec DPP subquery: " +
+              s"plan=${e.plan.getClass.getSimpleName}")
           e.updateResult()
         case _ =>
       }

spark/src/main/scala/org/apache/spark/sql/comet/CometSubqueryBroadcastExec.scala

@@ -40,14 +40,17 @@ import org.apache.spark.util.io.ChunkedByteBuffer
  * Comet replacement for SubqueryBroadcastExec that consumes Arrow broadcast data from a
  * CometBroadcastExchangeExec instead of HashedRelation from BroadcastExchangeExec.
  *
- * This enables broadcast exchange reuse between DPP subqueries and broadcast hash joins
- * when CometExecRule converts BroadcastExchangeExec to CometBroadcastExchangeExec.
- * Without this, the two exchanges have different types and canonical forms, so Spark's
- * ReuseExchangeAndSubquery (which runs after Comet rules) cannot match them.
+ * This enables broadcast exchange reuse between DPP subqueries and broadcast hash joins when
+ * CometExecRule converts BroadcastExchangeExec to CometBroadcastExchangeExec. Without this, the
+ * two exchanges have different types and canonical forms, so Spark's ReuseExchangeAndSubquery
+ * (which runs after Comet rules) cannot match them.
  *
- * @param indices the indices of the join keys in the list of keys from the build side
- * @param buildKeys the join keys from the build side of the join
- * @param child the CometBroadcastExchangeExec (or ReusedExchangeExec after reuse)
+ * @param indices
+ *   the indices of the join keys in the list of keys from the build side
+ * @param buildKeys
+ *   the join keys from the build side of the join
+ * @param child
+ *   the CometBroadcastExchangeExec (or ReusedExchangeExec after reuse)
  */
 case class CometSubqueryBroadcastExec(
     name: String,
@@ -150,8 +153,7 @@ case class CometSubqueryBroadcastExec(
 
   override def stringArgs: Iterator[Any] = super.stringArgs ++ Iterator(s"[id=#$id]")
 
-  override protected def withNewChildInternal(
-      newChild: SparkPlan): CometSubqueryBroadcastExec =
+  override protected def withNewChildInternal(newChild: SparkPlan): CometSubqueryBroadcastExec =
     copy(child = newChild)
 }
 

spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala

@@ -205,8 +205,7 @@ class CometExecSuite extends CometTestBase {
         val df = spark.sql(
           "select * from dpp_fact_bhj join dpp_dim_bhj on fact_date = dim_date where dim_id > 7")
         // Exclude ReusedExchangeExec — it appears inside the DPP subquery after exchange reuse
-        val (_, cometPlan) = checkSparkAnswerAndOperator(
-          df, classOf[ReusedExchangeExec])
+        val (_, cometPlan) = checkSparkAnswerAndOperator(df, classOf[ReusedExchangeExec])
 
         val nativeScans = cometPlan.collect { case s: CometNativeScanExec => s }
         assert(nativeScans.nonEmpty, "Expected CometNativeScanExec in plan")
@@ -265,17 +264,25 @@ class CometExecSuite extends CometTestBase {
     }
   }
 
-  test("DPP broadcast exchange reuse investigation") {
+  test("non-AQE DPP with BHJ reuses broadcast exchange") {
     withTempDir { dir =>
       val path = s"${dir.getAbsolutePath}/data"
       withSQLConf(CometConf.COMET_EXEC_ENABLED.key -> "false") {
-        spark.range(100).selectExpr(
-          "id % 10 as store_id", "cast(id * 2 as int) as date_id",
-          "cast(id * 3 as int) as product_id", "cast(id as int) as units_sold")
-          .write.partitionBy("store_id").parquet(s"$path/fact")
-        spark.range(10).selectExpr(
-          "cast(id as int) as store_id", "cast(id as string) as country")
-          .write.parquet(s"$path/dim")
+        spark
+          .range(100)
+          .selectExpr(
+            "id % 10 as store_id",
+            "cast(id * 2 as int) as date_id",
+            "cast(id * 3 as int) as product_id",
+            "cast(id as int) as units_sold")
+          .write
+          .partitionBy("store_id")
+          .parquet(s"$path/fact")
+        spark
+          .range(10)
+          .selectExpr("cast(id as int) as store_id", "cast(id as string) as country")
+          .write
+          .parquet(s"$path/dim")
       }
 
       withSQLConf(
@@ -285,53 +292,35 @@ class CometExecSuite extends CometTestBase {
         spark.read.parquet(s"$path/fact").createOrReplaceTempView("fact_reuse")
         spark.read.parquet(s"$path/dim").createOrReplaceTempView("dim_reuse")
 
-        val df = spark.sql(
-          """SELECT f.date_id, f.store_id
+        val df = spark.sql("""SELECT f.date_id, f.store_id
             |FROM fact_reuse f JOIN dim_reuse d
             |ON f.store_id = d.store_id
             |WHERE d.country = 'DE'""".stripMargin)
-        df.collect()
-        val plan = df.queryExecution.executedPlan
-        // scalastyle:off println
-        println(s"[REUSE-DEBUG] Plan:\n${plan.treeString}")
-
-        // Walk into subquery expressions to see what's inside
-        plan.foreach { node =>
-          node.expressions.foreach { expr =>
-            expr.foreach {
-              case sub: InSubqueryExec =>
-                println(s"[REUSE-DEBUG] Found InSubqueryExec in ${node.getClass.getSimpleName}")
-                println(s"[REUSE-DEBUG]   sub.plan class: ${sub.plan.getClass.getSimpleName}")
-                println(s"[REUSE-DEBUG]   sub.plan tree:\n${sub.plan.treeString}")
-                sub.plan match {
-                  case sb: SubqueryBroadcastExec =>
-                    println(s"[REUSE-DEBUG]   SubqueryBroadcast child: " +
-                      s"${sb.child.getClass.getSimpleName}")
-                    println(s"[REUSE-DEBUG]   SubqueryBroadcast child tree:\n" +
-                      s"${sb.child.treeString}")
-                  case other =>
-                    println(s"[REUSE-DEBUG]   sub.plan is: ${other.getClass.getSimpleName}")
-                }
-              case _ =>
-            }
-          }
-        }
+        val (_, cometPlan) = checkSparkAnswer(df)
 
-        val reused = collectWithSubqueries(plan) {
-          case e: ReusedExchangeExec => e
+        // DPP subquery should use CometSubqueryBroadcastExec (not SubqueryBroadcastExec)
+        val cometSubqueries = collectWithSubqueries(cometPlan) {
+          case s: CometSubqueryBroadcastExec => s
         }
-        println(s"[REUSE-DEBUG] ReusedExchangeExec count: ${reused.size}")
+        assert(
+          cometSubqueries.nonEmpty,
+          "Expected CometSubqueryBroadcastExec in plan for exchange reuse")
 
-        val broadcasts = collectWithSubqueries(plan) {
-          case e: BroadcastExchangeExec => ("BroadcastExchangeExec", e: SparkPlan)
-          case e: CometBroadcastExchangeExec => ("CometBroadcastExchangeExec", e: SparkPlan)
+        // Broadcast exchange should be reused — only one CometBroadcastExchangeExec,
+        // the other replaced by ReusedExchangeExec
+        val reused = collectWithSubqueries(cometPlan) { case e: ReusedExchangeExec =>
+          e
         }
-        println(s"[REUSE-DEBUG] Broadcast exchange count: ${broadcasts.size}")
-        broadcasts.foreach { case (typ, e) =>
-          println(s"[REUSE-DEBUG]   $typ hash=${e.canonicalized.hashCode()}")
-          println(s"[REUSE-DEBUG]   $typ child: ${e.children.map(_.getClass.getSimpleName)}")
+        assert(
+          reused.nonEmpty,
+          s"Expected ReusedExchangeExec for broadcast exchange reuse:\n${cometPlan.treeString}")
+
+        val broadcasts = collectWithSubqueries(cometPlan) { case e: CometBroadcastExchangeExec =>
+          e
         }
-        // scalastyle:on println
+        assert(
+          broadcasts.size == 1,
+          s"Expected exactly 1 CometBroadcastExchangeExec (other reused):\n${cometPlan.treeString}")
       }
     }
   }