fix refcount when materializing device buffer from GDS (NVIDIA#2095)

* fix refcount when materializing device buffer from GDS Signed-off-by: Rong Ou <rong.ou@gmail.com> * really fix it Signed-off-by: Rong Ou <rong.ou@gmail.com>
nartal1 · Apr 8, 2021 · 73de638 · 73de638
1 parent 328b825
commit 73de638
Showing 1 changed file with 7 additions and 5 deletions.
diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsBufferStore.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsBufferStore.scala
@@ -375,12 +375,14 @@ abstract class RapidsBufferStore(
         }
         throw new IllegalStateException(s"Unable to get device memory buffer for ID: $id")
       } else {
-        withResource(materializeMemoryBuffer) {
+        materializeMemoryBuffer match {
           case h: HostMemoryBuffer =>
-            closeOnExcept(DeviceMemoryBuffer.allocate(size)) { deviceBuffer =>
-              logDebug(s"copying from host $h to device $deviceBuffer")
-              deviceBuffer.copyFromHostBuffer(h)
-              deviceBuffer
+            withResource(h) { _ =>
+              closeOnExcept(DeviceMemoryBuffer.allocate(size)) { deviceBuffer =>
+                logDebug(s"copying from host $h to device $deviceBuffer")
+                deviceBuffer.copyFromHostBuffer(h)
+                deviceBuffer
+              }
             }
           case d: DeviceMemoryBuffer => d
           case b => throw new IllegalStateException(s"Unrecognized buffer: $b")