diff --git a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionFilesSorter.java b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionFilesSorter.java
index 86315d32950..2cc6c329280 100644
--- a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionFilesSorter.java
+++ b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionFilesSorter.java
@@ -459,7 +459,10 @@ public void initSortedShuffleFiles(String shuffleKey) {
 
   @VisibleForTesting
   public void updateSortedShuffleFiles(String shuffleKey, String fileId, long fileLength) {
-    sortedShuffleFiles.get(shuffleKey).add(fileId);
+    Set<String> shuffleFiles = sortedShuffleFiles.get(shuffleKey);
+    if (shuffleFiles != null) {
+      shuffleFiles.add(fileId);
+    }
     sortedFileCount.incrementAndGet();
     sortedFilesSize.addAndGet(fileLength);
   }