[Spark 4] Clear the watermark of keys without timers or state in the streaming stateful ParDo

runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/BeamStatefulProcessor.java

@@ -120,13 +120,22 @@ public void init(OutputMode outputMode, TimeMode timeMode) {
   }
 
   private void storeWatermark(SparkTimerInternals timerInternals) {
-    checkStateNotNull(beamWatermark).update(timerInternals.currentInputWatermarkTime().getMillis());
+    ValueState<Long> watermarkState = checkStateNotNull(beamWatermark);
+    // A key without timers and state keeps no row, its next visit is a first visit.
+    if (timerInternals.getTimers().isEmpty() && !checkStateNotNull(beamState).exists()) {
+      if (watermarkState.exists()) {
+        watermarkState.clear();
+      }
+    } else {
+      watermarkState.update(timerInternals.currentInputWatermarkTime().getMillis());
+    }
   }
 
   private SparkTimerInternals timerInternals(TimerValues timerValues) {
     ValueState<Long> watermarkState = checkStateNotNull(beamWatermark);
     // Late data is filtered against the watermark of the key's previous visit.
-    // A key seen for the first time or idle for several batches filters against an older watermark.
+    // A key seen for the first time, cleared after its last window expired or idle for several
+    // batches filters against an older watermark.
     Instant low =
         watermarkState.exists()
             ? new Instant(watermarkState.get())

runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTest.java

@@ -21,6 +21,7 @@
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 
+import java.io.File;
 import java.io.Serializable;
 import java.util.Arrays;
 import java.util.Set;
@@ -72,13 +73,15 @@ public class StatefulParDoStreamingTest implements Serializable {
   private static final String TAG_RESTART = "stateful-restart";
   private static final String TAG_CHAINED = "stateful-chained";
   private static final String TAG_NULL_KEY = "stateful-null-key";
+  private static final String TAG_EXPIRY = "stateful-expiry";
 
   @After
   public void tearDown() {
     TestUnboundedSource.forget(TAG_BATCHES);
     TestUnboundedSource.forget(TAG_RESTART);
     TestUnboundedSource.forget(TAG_CHAINED);
     TestUnboundedSource.forget(TAG_NULL_KEY);
+    TestUnboundedSource.forget(TAG_EXPIRY);
   }
 
   private static Pipeline pipeline(
@@ -190,6 +193,42 @@ public void testNullKey() throws Exception {
     assertTrue(collected.contains(5));
   }
 
+  @Test
+  public void testExpiredKeyLeavesNoState() throws Exception {
+    String checkpoint = tempFolder.newFolder("cp").getAbsolutePath();
+    String collectorId = StreamingTestUtils.newCollectorId(TAG_EXPIRY);
+    SparkStructuredStreamingPipelineOptions options =
+        StreamingTestUtils.streamingOptions(checkpoint);
+    options.setMaxRecordsPerBatch(1L);
+    Pipeline pipeline = Pipeline.create(options);
+    pipeline
+        .apply("Read", Read.from(new TestUnboundedSource(TAG_EXPIRY, 1, 8)))
+        .apply("KeySelector", ParDo.of(new ExpiryKeySelectorFn()))
+        .setCoder(KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()))
+        .apply("Window", Window.into(FixedWindows.of(Duration.millis(1000))))
+        .apply("StatefulParDo", ParDo.of(new DownstreamFn()))
+        .apply("Collect", ParDo.of(new StreamingTestUtils.CollectDoFn<>(collectorId)));
+
+    assertEquals(PipelineResult.State.DONE, StreamingTestUtils.run(pipeline).getState());
+    Set<String> collected = StreamingTestUtils.collected(collectorId);
+    assertTrue(collected.contains(TestUnboundedSource.element(TAG_EXPIRY, 0, 0)));
+    for (String stateVarName : Arrays.asList("beamState", "beamTimers", "beamWatermark")) {
+      assertEquals(stateVarName, 0L, stateRows(checkpoint, stateVarName));
+    }
+  }
+
+  /** Rows of a state variable of the query of the only leaf, checkpointed under 0. */
+  private static long stateRows(String checkpoint, String stateVarName) {
+    return SESSION
+        .getSession()
+        .read()
+        .format("statestore")
+        .option("path", new File(checkpoint, "0").getAbsolutePath())
+        .option("stateVarName", stateVarName)
+        .load()
+        .count();
+  }
+
   private static final class KeySelectorFn extends DoFn<String, KV<String, String>> {
     @ProcessElement
     public void process(@Element String element, OutputReceiver<KV<String, String>> out) {
@@ -311,6 +350,16 @@ public void onTimer(@Timestamp Instant ts, OutputReceiver<KV<String, String>> ou
     }
   }
 
+  private static final class ExpiryKeySelectorFn extends DoFn<String, KV<String, String>> {
+    @ProcessElement
+    public void process(@Element String element, OutputReceiver<KV<String, String>> out) {
+      int i = TestUnboundedSource.indexOf(element);
+      if (i == 0) {
+        out.output(KV.of("k1", element));
+      }
+    }
+  }
+
   private static final class DownstreamFn extends DoFn<KV<String, String>, String> {
     @StateId("state")
     private final StateSpec<ValueState<Integer>> stateSpec = StateSpecs.value();