diff --git a/gluten-celeborn/common/src/main/java/org/apache/spark/shuffle/gluten/celeborn/CelebornShuffleManager.java b/gluten-celeborn/common/src/main/java/org/apache/spark/shuffle/gluten/celeborn/CelebornShuffleManager.java
index 63fb0cc1b9bd..e61aeb3d7088 100644
--- a/gluten-celeborn/common/src/main/java/org/apache/spark/shuffle/gluten/celeborn/CelebornShuffleManager.java
+++ b/gluten-celeborn/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) {