diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/StateInternalsTest.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/StateInternalsTest.java
index e15249969f26..6988a492736a 100644
--- a/runners/core-java/src/test/java/org/apache/beam/runners/core/StateInternalsTest.java
+++ b/runners/core-java/src/test/java/org/apache/beam/runners/core/StateInternalsTest.java
@@ -225,6 +225,11 @@ public void testSet() throws Exception {
     assertThat(later.read(), hasItems("C", "D"));
     assertFalse(later.contains("A").read());
 
+    value.remove("B");
+    value.remove("C");
+    value.remove("D");
+    assertTrue(value.isEmpty().read());
+
     // clear
     value.clear();
     assertThat(value.read(), Matchers.emptyIterable());
@@ -389,6 +394,11 @@ public void testMap() throws Exception {
     // isEmpty
     assertFalse(value.isEmpty().read());
 
+    value.remove("B");
+    value.remove("D");
+    value.remove("E");
+    assertTrue(value.isEmpty().read());
+
     // clear
     value.clear();
     assertThat(value.entries().read(), Matchers.emptyIterable());
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/stateful/SparkStateInternals.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/stateful/SparkStateInternals.java
index 51ceb4c87309..4f744ab3ab12 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/stateful/SparkStateInternals.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/stateful/SparkStateInternals.java
@@ -432,7 +432,11 @@ private Map<MapKeyT, MapValueT> readAsMap() {
     public void remove(MapKeyT key) {
       Map<MapKeyT, MapValueT> sparkMapState = readAsMap();
       sparkMapState.remove(key);
-      writeValue(sparkMapState);
+      if (sparkMapState.isEmpty()) {
+        clear();
+      } else {
+        writeValue(sparkMapState);
+      }
     }
 
     @Override
@@ -537,7 +541,11 @@ public ReadableState<Boolean> addIfAbsent(InputT input) {
     public void remove(InputT input) {
       Set<InputT> sparkSetState = readAsSet();
       sparkSetState.remove(input);
-      writeValue(sparkSetState);
+      if (sparkSetState.isEmpty()) {
+        clear();
+      } else {
+        writeValue(sparkSetState);
+      }
     }
 
     @Override