Remove unnecessary toBeReturned field from serialized batch iterators

sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuColumnarBatchSerializer.scala

@@ -48,15 +48,13 @@ trait BaseSerializedTableIterator extends Iterator[(Int, ColumnarBatch)] {
 class SerializedBatchIterator(dIn: DataInputStream)
   extends BaseSerializedTableIterator {
   private[this] var nextHeader: Option[SerializedTableHeader] = None
-  private[this] var toBeReturned: Option[ColumnarBatch] = None
   private[this] var streamClosed: Boolean = false
 
   // Don't install the callback if in a unit test
   Option(TaskContext.get()).foreach { tc =>
     onTaskCompletion(tc) {
-      toBeReturned.foreach(_.close())
-      toBeReturned = None
       dIn.close()
+      streamClosed = true
     }
   }
 
@@ -80,23 +78,20 @@ class SerializedBatchIterator(dIn: DataInputStream)
     }
   }
 
-  private def tryReadNext(): Option[ColumnarBatch] = {
-    if (nextHeader.isEmpty) {
-      None
-    } else {
-      withResource(new NvtxRange("Read Batch", NvtxColor.YELLOW)) { _ =>
-        val header = nextHeader.get
-        if (header.getNumColumns > 0) {
-          // This buffer will later be concatenated into another host buffer before being
-          // sent to the GPU, so no need to use pinned memory for these buffers.
-          closeOnExcept(
-            HostMemoryBuffer.allocate(header.getDataLen, false)) { hostBuffer =>
-            JCudfSerialization.readTableIntoBuffer(dIn, header, hostBuffer)
-            Some(SerializedTableColumn.from(header, hostBuffer))
-          }
-        } else {
-          Some(SerializedTableColumn.from(header))
+  private def readNextBatch(): ColumnarBatch = {
+    withResource(new NvtxRange("Read Batch", NvtxColor.YELLOW)) { _ =>
+      val header = nextHeader.get
+      nextHeader = None
+      if (header.getNumColumns > 0) {
+        // This buffer will later be concatenated into another host buffer before being
+        // sent to the GPU, so no need to use pinned memory for these buffers.
+        closeOnExcept(
+          HostMemoryBuffer.allocate(header.getDataLen, false)) { hostBuffer =>
+          JCudfSerialization.readTableIntoBuffer(dIn, header, hostBuffer)
+          SerializedTableColumn.from(header, hostBuffer)
         }
+      } else {
+        SerializedTableColumn.from(header)
       }
     }
   }
@@ -107,17 +102,10 @@ class SerializedBatchIterator(dIn: DataInputStream)
   }
 
   override def next(): (Int, ColumnarBatch) = {
-    if (toBeReturned.isEmpty) {
-      peekNextBatchSize()
-      toBeReturned = tryReadNext()
-      if (nextHeader.isEmpty || toBeReturned.isEmpty) {
-        throw new NoSuchElementException("Walked off of the end...")
-      }
+    if (!hasNext) {
+      throw new NoSuchElementException("Walked off of the end...")
     }
-    val ret = toBeReturned.get
-    toBeReturned = None
-    nextHeader = None
-    (0, ret)
+    (0, readNextBatch())
   }
 }
 
@@ -498,15 +486,13 @@ object KudoSerializedTableColumn {
 class KudoSerializedBatchIterator(dIn: DataInputStream)
   extends BaseSerializedTableIterator {
   private[this] var nextHeader: Option[KudoTableHeader] = None
-  private[this] var toBeReturned: Option[ColumnarBatch] = None
   private[this] var streamClosed: Boolean = false
 
   // Don't install the callback if in a unit test
   Option(TaskContext.get()).foreach { tc =>
     onTaskCompletion(tc) {
-      toBeReturned.foreach(_.close())
-      toBeReturned = None
       dIn.close()
+      streamClosed = true
     }
   }
 
@@ -530,23 +516,20 @@ class KudoSerializedBatchIterator(dIn: DataInputStream)
     }
   }
 
-  private def tryReadNext(): Option[ColumnarBatch] = {
-    if (nextHeader.isEmpty) {
-      None
-    } else {
-      withResource(new NvtxRange("Read Batch", NvtxColor.YELLOW)) { _ =>
-        val header = nextHeader.get
-        if (header.getNumColumns > 0) {
-          // This buffer will later be concatenated into another host buffer before being
-          // sent to the GPU, so no need to use pinned memory for these buffers.
-          closeOnExcept(HostMemoryBuffer.allocate(header.getTotalDataLen, false)) { hostBuffer =>
-            hostBuffer.copyFromStream(0, dIn, header.getTotalDataLen)
-            val kudoTable = new KudoTable(header, hostBuffer)
-            Some(KudoSerializedTableColumn.from(kudoTable))
-          }
-        } else {
-          Some(KudoSerializedTableColumn.from(new KudoTable(header, null)))
+  private def readNextBatch(): ColumnarBatch = {
+    withResource(new NvtxRange("Read Batch", NvtxColor.YELLOW)) { _ =>
+      val header = nextHeader.get
+      nextHeader = None
+      if (header.getNumColumns > 0) {
+        // This buffer will later be concatenated into another host buffer before being
+        // sent to the GPU, so no need to use pinned memory for these buffers.
+        closeOnExcept(HostMemoryBuffer.allocate(header.getTotalDataLen, false)) { hostBuffer =>
+          hostBuffer.copyFromStream(0, dIn, header.getTotalDataLen)
+          val kudoTable = new KudoTable(header, hostBuffer)
+          KudoSerializedTableColumn.from(kudoTable)
         }
+      } else {
+        KudoSerializedTableColumn.from(new KudoTable(header, null))
       }
     }
   }
@@ -557,16 +540,9 @@ class KudoSerializedBatchIterator(dIn: DataInputStream)
   }
 
   override def next(): (Int, ColumnarBatch) = {
-    if (toBeReturned.isEmpty) {
-      peekNextBatchSize()
-      toBeReturned = tryReadNext()
-      if (nextHeader.isEmpty || toBeReturned.isEmpty) {
-        throw new NoSuchElementException("Walked off of the end...")
-      }
+    if (!hasNext) {
+      throw new NoSuchElementException("Walked off of the end...")
     }
-    val ret = toBeReturned.get
-    toBeReturned = None
-    nextHeader = None
-    (0, ret)
+    (0, readNextBatch())
   }
 }