fix: Avoid null exception in removeSubquery (#147)

In `CometScalarSubquery.removeSubquery`, we should make sure the `planId` is still in `subqueryMap` before we go to delete subquery from it. Otherwise, null exception could be thrown.
apache · Mar 1, 2024 · 5905151 · 5905151
1 parent 9aa42cd
commit 5905151
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);
+      }
     }
   }