Apply review comment

NVIDIA · Feb 4, 2022 · 474aef3 · 474aef3
1 parent 9aefdaf
commit 474aef3
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuCoalesceBatches.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuCoalesceBatches.scala
@@ -459,7 +459,7 @@ class GpuCoalesceIterator(iter: Iterator[ColumnarBatch],
       spillCallback))
 
   protected def popAll(): Array[ColumnarBatch] = {
-    closeOnExcept(batches.map(_.getColumnarBatch())) { wip =>
+    closeOnExcept(batches.safeMap(_.getColumnarBatch())) { wip =>
       batches.safeClose()
       batches.clear()
       wip.toArray
@@ -536,7 +536,7 @@ class GpuCompressionAwareCoalesceIterator(
   private[this] var codec: TableCompressionCodec = _
 
   override protected def popAll(): Array[ColumnarBatch] = {
-    closeOnExcept(batches.map(_.getColumnarBatch())) { wip =>
+    closeOnExcept(batches.safeMap(_.getColumnarBatch())) { wip =>
       batches.safeClose()
       batches.clear()