andygrove commented on code in PR #3266:
URL: https://github.com/apache/datafusion-comet/pull/3266#discussion_r2725614700
##########
spark/src/main/scala/org/apache/spark/sql/comet/CometNativeColumnarToRowExec.scala:
##########
@@ -64,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()
Review Comment:
@wForget this addresses once of your comments from previous PR
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
