getBarchFromMeta needs to add to the reference count

NVIDIA · Nov 13, 2020 · d48a629 · d48a629
1 parent a0182e6
commit d48a629
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/MetaUtils.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/MetaUtils.scala
@@ -241,10 +241,10 @@ object MetaUtils extends Arm {
   def getBatchFromMeta(deviceBuffer: DeviceMemoryBuffer,
       meta: TableMeta,
       sparkTypes: Array[DataType]): ColumnarBatch = {
-    withResource(new Array[GpuColumnVector](meta.columnMetasLength())) { columns =>
+    closeOnExcept(new ArrayBuffer[GpuColumnVector](meta.columnMetasLength())) { columns =>
       val columnMeta = new ColumnMeta
       (0 until meta.columnMetasLength).foreach { i =>
-        columns(i) = makeColumn(deviceBuffer, meta.columnMetas(columnMeta, i), sparkTypes(i))
+        columns.append(makeColumn(deviceBuffer, meta.columnMetas(columnMeta, i), sparkTypes(i)))
       }
       new ColumnarBatch(columns.toArray, meta.rowCount.toInt)
     }