From 55ce793c13a9e22780ce4662141d7ad7409db68e Mon Sep 17 00:00:00 2001
From: Liang-Chi Hsieh <viirya@gmail.com>
Date: Thu, 29 Feb 2024 16:22:43 -0800
Subject: [PATCH] fix: Avoid null exception in removeSubquery

---
 .../org/apache/spark/sql/comet/CometScalarSubquery.java   | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/spark/src/main/java/org/apache/spark/sql/comet/CometScalarSubquery.java b/spark/src/main/java/org/apache/spark/sql/comet/CometScalarSubquery.java
index 98dc02542..29984ebb5 100644
--- a/spark/src/main/java/org/apache/spark/sql/comet/CometScalarSubquery.java
+++ b/spark/src/main/java/org/apache/spark/sql/comet/CometScalarSubquery.java
@@ -47,10 +47,12 @@ public static synchronized void setSubquery(long planId, ScalarSubquery subquery
   }
 
   public static synchronized void removeSubquery(long planId, ScalarSubquery subquery) {
-    subqueryMap.get(planId).remove(subquery.exprId().id());
+    if (subqueryMap.containsKey(planId)) {
+      subqueryMap.get(planId).remove(subquery.exprId().id());
 
-    if (subqueryMap.get(planId).isEmpty()) {
-      subqueryMap.remove(planId);
+      if (subqueryMap.get(planId).isEmpty()) {
+        subqueryMap.remove(planId);
+      }
     }
   }