Add doExecuteBroadcast support to CometNativeColumnarToRowExec

andygrove · claude · andygrove · commit 2581ea5402c0 · 2026-01-21T08:08:50.000-07:00
When native columnar-to-row conversion is enabled (now the default),
CometNativeColumnarToRowExec is used instead of CometColumnarToRowExec.
However, it was missing the doExecuteBroadcast implementation required
for broadcast exchange operations, causing test failures.

Changes:
- Add doExecuteBroadcast implementation to CometNativeColumnarToRowExec
  that uses the native converter for broadcast data transformation
- Update CometExecSuite test to handle both CometColumnarToRowExec and
  CometNativeColumnarToRowExec
- Fix parent-child relationship check to account for InputAdapter wrapper
  nodes used by Spark's codegen
- Remove nodeName override from CometNativeColumnarToRowExec

Co-Authored-By: Claude Opus 4.5 &lt;noreply@anthropic.com&gt;
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeColumnarToRowExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeColumnarToRowExec.scala
@@ -19,15 +19,25 @@
 
 package org.apache.spark.sql.comet
 
-import org.apache.spark.TaskContext
+import java.util.UUID
+import java.util.concurrent.{Future, TimeoutException, TimeUnit}
+
+import scala.concurrent.Promise
+import scala.util.control.NonFatal
+
+import org.apache.spark.{broadcast, SparkException, TaskContext}
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.{Attribute, SortOrder}
 import org.apache.spark.sql.catalyst.plans.physical.Partitioning
-import org.apache.spark.sql.execution.{ColumnarToRowTransition, SparkPlan}
+import org.apache.spark.sql.comet.util.{Utils => CometUtils}
+import org.apache.spark.sql.errors.QueryExecutionErrors
+import org.apache.spark.sql.execution.{ColumnarToRowTransition, SparkPlan, SQLExecution}
+import org.apache.spark.sql.execution.adaptive.BroadcastQueryStageExec
+import org.apache.spark.sql.execution.exchange.ReusedExchangeExec
 import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
 import org.apache.spark.sql.types.StructType
-import org.apache.spark.util.Utils
+import org.apache.spark.util.{SparkFatalException, Utils}
 
 import org.apache.comet.{CometConf, NativeColumnarToRowConverter}
 
@@ -53,9 +63,6 @@ case class CometNativeColumnarToRowExec(child: SparkPlan)
   // supportsColumnar requires to be only called on driver side, see also SPARK-37779.
   assert(Utils.isInRunningSparkTask || child.supportsColumnar)
 
-  // Use the same display name as CometColumnarToRowExec for plan compatibility
-  override def nodeName: String = "CometColumnarToRow"
-
   override def output: Seq[Attribute] = child.output
 
   override def outputPartitioning: Partitioning = child.outputPartitioning
@@ -67,6 +74,105 @@ case class CometNativeColumnarToRowExec(child: SparkPlan)
     "numInputBatches" -> SQLMetrics.createMetric(sparkContext, "number of input batches"),
     "convertTime" -> SQLMetrics.createNanoTimingMetric(sparkContext, "time in conversion"))
 
+  @transient
+  private lazy val promise = Promise[broadcast.Broadcast[Any]]()
+
+  @transient
+  private val timeout: Long = conf.broadcastTimeout
+
+  private val runId: UUID = UUID.randomUUID
+
+  private lazy val cometBroadcastExchange = findCometBroadcastExchange(child)
+
+  @transient
+  lazy val relationFuture: Future[broadcast.Broadcast[Any]] = {
+    SQLExecution.withThreadLocalCaptured[broadcast.Broadcast[Any]](
+      session,
+      CometBroadcastExchangeExec.executionContext) {
+      try {
+        // Setup a job group here so later it may get cancelled by groupId if necessary.
+        sparkContext.setJobGroup(
+          runId.toString,
+          s"CometNativeColumnarToRow broadcast exchange (runId $runId)",
+          interruptOnCancel = true)
+
+        val numOutputRows = longMetric("numOutputRows")
+        val numInputBatches = longMetric("numInputBatches")
+        val localSchema = this.schema
+        val batchSize = CometConf.COMET_BATCH_SIZE.get()
+        val broadcastColumnar = child.executeBroadcast()
+        val serializedBatches =
+          broadcastColumnar.value.asInstanceOf[Array[org.apache.spark.util.io.ChunkedByteBuffer]]
+
+        // Use native converter to convert columnar data to rows
+        val converter = new NativeColumnarToRowConverter(localSchema, batchSize)
+        try {
+          val rows = serializedBatches.iterator
+            .flatMap(CometUtils.decodeBatches(_, this.getClass.getSimpleName))
+            .flatMap { batch =>
+              numInputBatches += 1
+              numOutputRows += batch.numRows()
+              converter.convert(batch)
+            }
+
+          val mode = cometBroadcastExchange.get.mode
+          val relation = mode.transform(rows, Some(numOutputRows.value))
+          val broadcasted = sparkContext.broadcastInternal(relation, serializedOnly = true)
+          val executionId = sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY)
+          SQLMetrics.postDriverMetricUpdates(sparkContext, executionId, metrics.values.toSeq)
+          promise.trySuccess(broadcasted)
+          broadcasted
+        } finally {
+          converter.close()
+        }
+      } catch {
+        // SPARK-24294: To bypass scala bug: https://github.com/scala/bug/issues/9554, we throw
+        // SparkFatalException, which is a subclass of Exception. ThreadUtils.awaitResult
+        // will catch this exception and re-throw the wrapped fatal throwable.
+        case oe: OutOfMemoryError =>
+          val ex = new SparkFatalException(oe)
+          promise.tryFailure(ex)
+          throw ex
+        case e if !NonFatal(e) =>
+          val ex = new SparkFatalException(e)
+          promise.tryFailure(ex)
+          throw ex
+        case e: Throwable =>
+          promise.tryFailure(e)
+          throw e
+      }
+    }
+  }
+
+  override def doExecuteBroadcast[T](): broadcast.Broadcast[T] = {
+    if (cometBroadcastExchange.isEmpty) {
+      throw new SparkException(
+        "CometNativeColumnarToRowExec only supports doExecuteBroadcast when child contains a " +
+          "CometBroadcastExchange, but got " + child)
+    }
+
+    try {
+      relationFuture.get(timeout, TimeUnit.SECONDS).asInstanceOf[broadcast.Broadcast[T]]
+    } catch {
+      case ex: TimeoutException =>
+        logError(s"Could not execute broadcast in $timeout secs.", ex)
+        if (!relationFuture.isDone) {
+          sparkContext.cancelJobGroup(runId.toString)
+          relationFuture.cancel(true)
+        }
+        throw QueryExecutionErrors.executeBroadcastTimeoutError(timeout, Some(ex))
+    }
+  }
+
+  private def findCometBroadcastExchange(op: SparkPlan): Option[CometBroadcastExchangeExec] = {
+    op match {
+      case b: CometBroadcastExchangeExec => Some(b)
+      case b: BroadcastQueryStageExec => findCometBroadcastExchange(b.plan)
+      case b: ReusedExchangeExec => findCometBroadcastExchange(b.child)
+      case _ => op.children.collectFirst(Function.unlift(findCometBroadcastExchange))
+    }
+  }
+
   override def doExecute(): RDD[InternalRow] = {
     val numOutputRows = longMetric("numOutputRows")
     val numInputBatches = longMetric("numInputBatches")
diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala
@@ -35,7 +35,7 @@ import org.apache.spark.sql.catalyst.expressions.{Expression, ExpressionInfo, He
 import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateMode, BloomFilterAggregate}
 import org.apache.spark.sql.comet._
 import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometShuffleExchangeExec}
-import org.apache.spark.sql.execution.{CollectLimitExec, ProjectExec, SQLExecution, UnionExec}
+import org.apache.spark.sql.execution.{CollectLimitExec, ProjectExec, SparkPlan, SQLExecution, UnionExec}
 import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, BroadcastQueryStageExec}
 import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat
 import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ReusedExchangeExec, ShuffleExchangeExec}
@@ -864,9 +864,11 @@ class CometExecSuite extends CometTestBase {
             checkSparkAnswerAndOperator(df)
 
             // Before AQE: one CometBroadcastExchange, no CometColumnarToRow
-            var columnarToRowExec = stripAQEPlan(df.queryExecution.executedPlan).collect {
-              case s: CometColumnarToRowExec => s
-            }
+            var columnarToRowExec: Seq[SparkPlan] =
+              stripAQEPlan(df.queryExecution.executedPlan).collect {
+                case s: CometColumnarToRowExec => s
+                case s: CometNativeColumnarToRowExec => s
+              }
             assert(columnarToRowExec.isEmpty)
 
             // Disable CometExecRule after the initial plan is generated. The CometSortMergeJoin and
@@ -880,14 +882,25 @@ class CometExecSuite extends CometTestBase {
             // After AQE: CometBroadcastExchange has to be converted to rows to conform to Spark
             // BroadcastHashJoin.
             val plan = stripAQEPlan(df.queryExecution.executedPlan)
-            columnarToRowExec = plan.collect { case s: CometColumnarToRowExec =>
-              s
+            columnarToRowExec = plan.collect {
+              case s: CometColumnarToRowExec => s
+              case s: CometNativeColumnarToRowExec => s
             }
             assert(columnarToRowExec.length == 1)
 
-            // This ColumnarToRowExec should be the immediate child of BroadcastHashJoinExec
-            val parent = plan.find(_.children.contains(columnarToRowExec.head))
-            assert(parent.get.isInstanceOf[BroadcastHashJoinExec])
+            // This ColumnarToRowExec should be a descendant of BroadcastHashJoinExec (possibly
+            // wrapped by InputAdapter for codegen).
+            val broadcastJoins = plan.collect { case b: BroadcastHashJoinExec => b }
+            assert(broadcastJoins.nonEmpty, s"Expected BroadcastHashJoinExec in plan:\n$plan")
+            val hasC2RDescendant = broadcastJoins.exists { join =>
+              join.find {
+                case _: CometColumnarToRowExec | _: CometNativeColumnarToRowExec => true
+                case _ => false
+              }.isDefined
+            }
+            assert(
+              hasC2RDescendant,
+              "BroadcastHashJoinExec should have a columnar-to-row descendant")
 
             // There should be a CometBroadcastExchangeExec under CometColumnarToRowExec
             val broadcastQueryStage =
