diff --git a/helix-core/src/main/java/org/apache/helix/controller/stages/MessageGenerationPhase.java b/helix-core/src/main/java/org/apache/helix/controller/stages/MessageGenerationPhase.java
index 859c6679e9..722b617d3c 100644
--- a/helix-core/src/main/java/org/apache/helix/controller/stages/MessageGenerationPhase.java
+++ b/helix-core/src/main/java/org/apache/helix/controller/stages/MessageGenerationPhase.java
@@ -152,8 +152,12 @@ private void generateMessage(final Resource resource, final BaseControllerDataPr
       // resourceStateMap. This instance may not have had been dropped by the rebalance strategy.
       // This check is required to ensure that the instances removed from the ideal state stateMap
       // are properly dropped.
+      // This should only solve for instance operation case where the instance is removed from the statemap but there
+      // are still valid assignments in the mapping. We should not consider case where there is no mapping at all for
+      // the resource, which can occur on a rebalance failure. If the resource has been removed, the BP will
+      // contain the DROPPED states
       for (String instance : currentStateMap.keySet()) {
-        if (!instanceStateMap.containsKey(instance)) {
+        if (!instanceStateMap.isEmpty() && !instanceStateMap.containsKey(instance)) {
           instanceStateMap.put(instance, HelixDefinedState.DROPPED.name());
         }
       }