RapidsBufferStore race could cause NPE (NVIDIA#928)

Signed-off-by: Alessandro Bellina <abellina@nvidia.com>
sperlingxx · Nov 20, 2020 · d28013e · d28013e
1 parent aec54d2
commit d28013e
Showing 1 changed file with 3 additions and 1 deletion.
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
@@ -213,8 +213,10 @@ abstract class RapidsBufferStore(
     val bufferToSpill = buffers.nextSpillableBuffer()
     if (bufferToSpill != null) {
       spillAndFreeBuffer(bufferToSpill, stream)
+      bufferToSpill.size
+    } else {
+      0
     }
-    bufferToSpill.size
   }
 
   private def spillAndFreeBuffer(buffer: RapidsBufferBase, stream: Cuda.Stream): Unit = {