diff --git a/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java b/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
index 31677f7b8..b43bd9598 100644
--- a/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
+++ b/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
@@ -1888,6 +1888,8 @@ public class ShuffleClientImpl extends ShuffleClient {
       cause = StatusCode.PUSH_DATA_PRIMARY_WORKER_EXCLUDED;
     } else if (message.startsWith(StatusCode.PUSH_DATA_REPLICA_WORKER_EXCLUDED.name())) {
       cause = StatusCode.PUSH_DATA_REPLICA_WORKER_EXCLUDED;
+    } else if (message.startsWith(StatusCode.PUSH_DATA_FAIL_PARTITION_NOT_FOUND.name())) {
+      cause = StatusCode.PUSH_DATA_FAIL_PARTITION_NOT_FOUND;
     } else if (ExceptionUtils.connectFail(message)) {
       // Throw when push to primary worker connection causeException.
       cause = StatusCode.PUSH_DATA_CONNECTION_EXCEPTION_PRIMARY;