[CELEBORN] Avoid CelebornShuffleManager#getWriter adding shuffle id r…

…epeatedly to columnarShuffleIds (#6281)
apache · Jul 1, 2024 · f80c068 · f80c068
1 parent d516f56
commit f80c068
Showing 1 changed file with 2 additions and 7 deletions.
diff --git a/...common/src/main/java/org/apache/spark/shuffle/gluten/celeborn/CelebornShuffleManager.java b/...common/src/main/java/org/apache/spark/shuffle/gluten/celeborn/CelebornShuffleManager.java
@@ -215,12 +215,8 @@ public <K, V, C> ShuffleHandle registerShuffle(
 
   @Override
   public boolean unregisterShuffle(int shuffleId) {
-    if (columnarShuffleIds.contains(shuffleId)) {
-      if (columnarShuffleManager().unregisterShuffle(shuffleId)) {
-        return columnarShuffleIds.remove(shuffleId);
-      } else {
-        return false;
-      }
+    if (columnarShuffleIds.remove(shuffleId)) {
+      return columnarShuffleManager().unregisterShuffle(shuffleId);
     }
     return CelebornUtils.unregisterShuffle(
         lifecycleManager,
@@ -311,7 +307,6 @@ public <K, V> ShuffleWriter<K, V> getWriter(
           return vanillaCelebornShuffleManager().getWriter(handle, mapId, context, metrics);
         }
       } else {
-        columnarShuffleIds.add(handle.shuffleId());
         return columnarShuffleManager().getWriter(handle, mapId, context, metrics);
       }
     } catch (Exception e) {