fix: Avoid null exception in removeSubquery

viirya · Mar 1, 2024 · 55ce793 · 55ce793
1 parent 9aa42cd
commit 55ce793
Showing 1 changed file with 5 additions and 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
@@ -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);
+      }
     }
   }