Ensure we exhaust row iterator before closing the hostBatch

NVIDIA · Jun 27, 2023 · 48e8c05 · 48e8c05
1 parent 0535357
commit 48e8c05
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/execution/GpuBroadcastToRowExec.scala b/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/execution/GpuBroadcastToRowExec.scala
@@ -110,7 +110,7 @@ case class GpuBroadcastToRowExec(
               BoundReference(idx, buildKeys(idx).dataType, buildKeys(idx).nullable))
             rowProject(broadcastRow).copy().asInstanceOf[InternalRow]
         }
-      }.toArray
+      }.toArray // force evaluation so we don't close hostBatch too soon
     }
 
     gpuLongMetric("dataSize") += serBatch.dataSize

diff --git a/...lugin/src/main/scala/org/apache/spark/sql/rapids/execution/GpuSubqueryBroadcastExec.scala b/...lugin/src/main/scala/org/apache/spark/sql/rapids/execution/GpuSubqueryBroadcastExec.scala
@@ -260,13 +260,13 @@ case class GpuSubqueryBroadcastExec(
       hostBatch.rowIterator().asScala.map { row =>
         val broadcastRow = broadcastModeProject.map(_(row)).getOrElse(row)
         rowProject(broadcastRow).copy().asInstanceOf[InternalRow]
-      }
+      }.toArray // force evaluation so we don't close hostBatch too soon
     }
 
     gpuLongMetric("dataSize") += serBatch.dataSize
     gpuLongMetric(COLLECT_TIME) += System.nanoTime() - beforeCollect
 
-    result.toArray
+    result
   }
 
   protected override def doExecute(): RDD[InternalRow] = {