[Spark 4] Filter late data against the key's previous watermark in the streaming stateful ParDo

.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Spark.json

@@ -1,4 +1,4 @@
 {
     "comment": "Modify this file in a trivial way to cause this test suite to run",
-    "modification": 5
+    "modification": 6
 }

.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_SparkStructuredStreaming.json

@@ -11,5 +11,6 @@
   "https://github.com/apache/beam/pull/39793": "noting that PR #39793 should run this test",
   "https://github.com/apache/beam/pull/40103": "noting that PR #40103 should run this test",
   "https://github.com/apache/beam/issues/40427": "Spark ValidatesRunner tests run serially",
-  "https://github.com/apache/beam/pull/40281": "noting that PR #40281 should run this test"
+  "https://github.com/apache/beam/pull/40281": "noting that PR #40281 should run this test",
+  "https://github.com/apache/beam/issues/36841": "Spark 4 streaming watermark lag"
 }

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

@@ -72,6 +72,7 @@ public class BeamStatefulProcessor<K, V, OutputT>
 
   private static final String BEAM_STATE = "beamState";
   private static final String BEAM_TIMERS = "beamTimers";
+  private static final String BEAM_WATERMARK = "beamWatermark";
 
   private final DoFn<KV<K, V>, OutputT> doFn;
   private final DoFnRunnerFactory<KV<K, V>, OutputT> runnerFactory;
@@ -85,6 +86,7 @@ public class BeamStatefulProcessor<K, V, OutputT>
 
   private transient @Nullable MapState<String, byte[]> beamState;
   private transient @Nullable ValueState<byte[]> beamTimers;
+  private transient @Nullable ValueState<Long> beamWatermark;
   private transient @Nullable List<WindowedValue<OutputT>> currentOutputs;
   private transient boolean needsBundleStart;
 
@@ -113,45 +115,64 @@ public void init(OutputMode outputMode, TimeMode timeMode) {
     beamState =
         getHandle().getMapState(BEAM_STATE, Encoders.STRING(), Encoders.BINARY(), TTLConfig.NONE());
     beamTimers = getHandle().getValueState(BEAM_TIMERS, Encoders.BINARY(), TTLConfig.NONE());
+    beamWatermark = getHandle().getValueState(BEAM_WATERMARK, Encoders.LONG(), TTLConfig.NONE());
+  }
+
+  private void storeWatermark(SparkTimerInternals timerInternals) {
+    checkStateNotNull(beamWatermark).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.
+    Instant low =
+        watermarkState.exists()
+            ? new Instant(watermarkState.get())
+            : BoundedWindow.TIMESTAMP_MIN_VALUE;
+    Instant high = new Instant(timerValues.getCurrentWatermarkInMs());
+    SparkTimerInternals timerInternals = SparkTimerInternals.forWatermarks(low, high);
+    restoreTimers(timerInternals);
+    return timerInternals;
   }
 
   @Override
   public Iterator<WindowedValue<OutputT>> handleInputRows(
       K key, Iterator<WindowedValue<KV<K, V>>> rows, TimerValues timerValues) {
     List<WindowedValue<OutputT>> outputs = new ArrayList<>();
-    SparkTimerInternals timerInternals =
-        SparkTimerInternals.forWatermark(new Instant(timerValues.getCurrentWatermarkInMs()));
-    restoreTimers(timerInternals);
+    SparkTimerInternals timerInternals = timerInternals(timerValues);
 
     DoFnRunner<KV<K, V>, OutputT> runner = keyRunner(key, timerInternals, outputs);
     while (rows.hasNext()) {
       runner.processElement(rows.next());
     }
+    timerInternals.advanceWatermark();
     // The next bundle opens in keyRunner.
     runner.finishBundle();
     needsBundleStart = true;
 
     persistTimers(timerInternals);
     reconcileWakeupTimer(timerInternals, null);
+    storeWatermark(timerInternals);
     return ScalaInterop.scalaIterator(outputs);
   }
 
   @Override
   public Iterator<WindowedValue<OutputT>> handleExpiredTimer(
       K key, TimerValues timerValues, ExpiredTimerInfo expiredTimerInfo) {
     List<WindowedValue<OutputT>> outputs = new ArrayList<>();
-    Instant watermark = new Instant(timerValues.getCurrentWatermarkInMs());
-    SparkTimerInternals timerInternals = SparkTimerInternals.forWatermark(watermark);
-    restoreTimers(timerInternals);
+    SparkTimerInternals timerInternals = timerInternals(timerValues);
+    timerInternals.advanceWatermark();
 
     DoFnRunner<KV<K, V>, OutputT> runner = keyRunner(key, timerInternals, outputs);
-    fireDueTimers(key, watermark, timerInternals, runner);
+    fireDueTimers(key, timerInternals.currentInputWatermarkTime(), timerInternals, runner);
     // The next bundle opens in keyRunner.
     runner.finishBundle();
     needsBundleStart = true;
 
     persistTimers(timerInternals);
     reconcileWakeupTimer(timerInternals, expiredTimerInfo.getExpiryTimeInMs());
+    storeWatermark(timerInternals);
     return ScalaInterop.scalaIterator(outputs);
   }
 

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

@@ -43,6 +43,8 @@
 import org.apache.beam.sdk.state.ValueState;
 import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.windowing.FixedWindows;
+import org.apache.beam.sdk.transforms.windowing.Window;
 import org.apache.beam.sdk.values.KV;
 import org.joda.time.Duration;
 import org.joda.time.Instant;
@@ -66,11 +68,13 @@ public class StatefulParDoStreamingTest implements Serializable {
 
   private static final String TAG_BATCHES = "stateful-batches";
   private static final String TAG_RESTART = "stateful-restart";
+  private static final String TAG_CHAINED = "stateful-chained";
 
   @After
   public void tearDown() {
     TestUnboundedSource.forget(TAG_BATCHES);
     TestUnboundedSource.forget(TAG_RESTART);
+    TestUnboundedSource.forget(TAG_CHAINED);
   }
 
   private static Pipeline pipeline(
@@ -141,6 +145,28 @@ public void testCheckpointRestartPreservesStateAndTimers() throws Exception {
     assertFalse(collected.contains("k1-started"));
   }
 
+  @Test
+  public void testChainedStatefulParDoTimerOutputSurvives() throws Exception {
+    String collectorId = StreamingTestUtils.newCollectorId(TAG_CHAINED);
+    SparkStructuredStreamingPipelineOptions options =
+        StreamingTestUtils.streamingOptions(tempFolder);
+    options.setMaxRecordsPerBatch(1L);
+    Pipeline pipeline = Pipeline.create(options);
+    pipeline
+        .apply("Read", Read.from(new TestUnboundedSource(TAG_CHAINED, 1, 4)))
+        .apply("KeySelector", ParDo.of(new KeySelectorFn()))
+        .setCoder(KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()))
+        .apply("EmitTimer", ParDo.of(new TimerEmitterFn()))
+        .setCoder(KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()))
+        .apply("Window", Window.into(FixedWindows.of(Duration.millis(500))))
+        .apply("Downstream", 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("timer-output"));
+  }
+
   private static final class KeySelectorFn extends DoFn<String, KV<String, String>> {
     @ProcessElement
     public void process(@Element String element, OutputReceiver<KV<String, String>> out) {
@@ -239,4 +265,40 @@ public void onTimer(OutputReceiver<String> out) {
       out.output("k1-timer");
     }
   }
+
+  private static final class TimerEmitterFn extends DoFn<KV<String, String>, KV<String, String>> {
+    @TimerId("t")
+    private final TimerSpec timer = TimerSpecs.timer(TimeDomain.EVENT_TIME);
+
+    @ProcessElement
+    public void process(
+        @Element KV<String, String> element,
+        @Timestamp Instant ts,
+        @TimerId("t") Timer timer,
+        OutputReceiver<KV<String, String>> out) {
+      out.output(KV.of("downstream", element.getValue()));
+      if (TestUnboundedSource.indexOf(element.getValue()) == 0) {
+        timer.set(ts.plus(Duration.millis(250)));
+      }
+    }
+
+    @OnTimer("t")
+    public void onTimer(@Timestamp Instant ts, OutputReceiver<KV<String, String>> out) {
+      out.outputWithTimestamp(KV.of("downstream", "timer-output"), ts);
+    }
+  }
+
+  private static final class DownstreamFn extends DoFn<KV<String, String>, String> {
+    @StateId("state")
+    private final StateSpec<ValueState<Integer>> stateSpec = StateSpecs.value();
+
+    @ProcessElement
+    public void process(
+        @Element KV<String, String> element,
+        @StateId("state") ValueState<Integer> state,
+        OutputReceiver<String> out) {
+      state.write(1);
+      out.output(element.getValue());
+    }
+  }
 }

runners/spark/src/main/java/org/apache/beam/runners/spark/stateful/SparkTimerInternals.java

@@ -54,9 +54,9 @@ private SparkTimerInternals(
     this.synchronizedProcessingTime = synchronizedProcessingTime;
   }
 
-  /** Build a {@link TimerInternals} initialized with a given watermark. */
-  public static SparkTimerInternals forWatermark(Instant watermark) {
-    return new SparkTimerInternals(watermark, watermark, new Instant(0));
+  /** Build a {@link TimerInternals} initialized with low and high watermarks. */
+  public static SparkTimerInternals forWatermarks(Instant low, Instant high) {
+    return new SparkTimerInternals(low, high, new Instant(0));
   }
 
   /** Build the {@link TimerInternals} according to the feeding streams. */

runners/spark/src/test/java/org/apache/beam/runners/spark/stateful/SparkTimerInternalsTest.java

@@ -103,4 +103,14 @@ public void testGetNextProcessingTimerIgnoresEventTimeTimers() {
 
     assertNull(timerInternals.getNextProcessingTimer());
   }
+
+  @Test
+  public void testForWatermarksAdvancesFromLowToHigh() {
+    SparkTimerInternals timerInternals =
+        SparkTimerInternals.forWatermarks(new Instant(1000), new Instant(2000));
+    assertEquals(new Instant(1000), timerInternals.currentInputWatermarkTime());
+
+    timerInternals.advanceWatermark();
+    assertEquals(new Instant(2000), timerInternals.currentInputWatermarkTime());
+  }
 }