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);
+      }
     }
   }