Copilot commented on code in PR #3266:
URL: https://github.com/apache/datafusion-comet/pull/3266#discussion_r2729909473
##########
spark/src/main/scala/org/apache/spark/sql/comet/CometNativeColumnarToRowExec.scala:
##########
@@ -64,6 +74,116 @@ 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()
+ val result = converter.convert(batch)
+ // Wrap iterator to close batch after consumption
+ new Iterator[InternalRow] {
+ override def hasNext: Boolean = {
+ val hasMore = result.hasNext
+ if (!hasMore) {
+ batch.close()
+ }
+ hasMore
+ }
+ override def next(): InternalRow = result.next()
+ }
Review Comment:
`batch.close()` is only triggered when the converted row iterator is fully
exhausted. If the consumer short-circuits (e.g., LIMIT / take / early
termination) or if `converter.convert(batch)` throws, the `ColumnarBatch` may
never be closed, leaking off-heap/native memory. Consider ensuring the batch is
closed on task completion (e.g., via `TaskContext.addTaskCompletionListener`)
and closing the batch on conversion failure (try/catch around
`converter.convert`).
##########
spark/src/main/scala/org/apache/spark/sql/comet/CometNativeColumnarToRowExec.scala:
##########
@@ -91,7 +211,17 @@ case class CometNativeColumnarToRowExec(child: SparkPlan)
val result = converter.convert(batch)
convertTime += System.nanoTime() - startTime
- result
+ // Wrap iterator to close batch after consumption
+ new Iterator[InternalRow] {
+ override def hasNext: Boolean = {
+ val hasMore = result.hasNext
+ if (!hasMore) {
+ batch.close()
+ }
+ hasMore
+ }
+ override def next(): InternalRow = result.next()
+ }
Review Comment:
In `doExecute`, the `ColumnarBatch` is only closed when `hasNext` returns
false. If downstream stops consuming rows early, the batch won't be closed
(resource leak). Consider using a task completion listener (or a completion
iterator) to always close the current batch even when the iterator is not fully
consumed.
--
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]
