WAR for a possible npe issue

Signed-off-by: Firestarman <[email protected]>
nvliyuan · May 15, 2024 · 42fed9f · 42fed9f
1 parent 9ce9b12
commit 42fed9f
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuColumnarBatchSerializer.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuColumnarBatchSerializer.scala
@@ -146,7 +146,9 @@ private class GpuColumnarBatchSerializerInstance(dataSize: GpuMetric, serTime: G
       with Logging {
     private[this] val dOut = new DataOutputStream(new BufferedOutputStream(out))
     private[this] val tableSerializer = new SimpleTableSerializer()
-    onTaskCompletion(TaskContext.get())(tableSerializer.close())
+    Option(TaskContext.get()).foreach { tc =>
+      onTaskCompletion(tc)(tableSerializer.close())
+    }
 
     private lazy val serializeBatch: ColumnarBatch => Unit = if (isSerializedTable) {
       serializeGpuBatch