From 3f593fb652a21c3e35c547f94f5241a1f1004326 Mon Sep 17 00:00:00 2001 From: Tobias Kaymak Date: Wed, 7 Oct 2026 11:23:39 +0000 Subject: [PATCH 1/3] [Spark 4] Filter late data against the key's previous watermark in the streaming stateful ParDo Late data is filtered against the watermark of the key's previous visit and timers fire against the current one, so timer output behind the shared watermark survives into a downstream stateful ParDo. Part of #36841. --- ...PostCommit_Java_ValidatesRunner_Spark.json | 2 +- ...idatesRunner_SparkStructuredStreaming.json | 3 +- .../state/BeamStatefulProcessor.java | 35 ++++++++--- .../streaming/StatefulParDoStreamingTest.java | 62 +++++++++++++++++++ .../spark/stateful/SparkTimerInternals.java | 6 +- 5 files changed, 96 insertions(+), 12 deletions(-) diff --git a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Spark.json b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Spark.json index 2160d3c68005..a84f69a97721 100644 --- a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Spark.json +++ b/.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 } diff --git a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_SparkStructuredStreaming.json b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_SparkStructuredStreaming.json index 2b28539c14fa..bd93fd4e9881 100644 --- a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_SparkStructuredStreaming.json +++ b/.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" } diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/BeamStatefulProcessor.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/BeamStatefulProcessor.java index c307ae6f37ba..e359ed6bd58b 100644 --- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/BeamStatefulProcessor.java +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/BeamStatefulProcessor.java @@ -72,6 +72,7 @@ public class BeamStatefulProcessor private static final String BEAM_STATE = "beamState"; private static final String BEAM_TIMERS = "beamTimers"; + private static final String BEAM_WATERMARK = "beamWatermark"; private final DoFn, OutputT> doFn; private final DoFnRunnerFactory, OutputT> runnerFactory; @@ -85,6 +86,7 @@ public class BeamStatefulProcessor private transient @Nullable MapState beamState; private transient @Nullable ValueState beamTimers; + private transient @Nullable ValueState beamWatermark; private transient @Nullable List> currentOutputs; private transient boolean needsBundleStart; @@ -113,26 +115,45 @@ 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 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> handleInputRows( K key, Iterator>> rows, TimerValues timerValues) { List> outputs = new ArrayList<>(); - SparkTimerInternals timerInternals = - SparkTimerInternals.forWatermark(new Instant(timerValues.getCurrentWatermarkInMs())); - restoreTimers(timerInternals); + SparkTimerInternals timerInternals = timerInternals(timerValues); DoFnRunner, 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); } @@ -140,18 +161,18 @@ public Iterator> handleInputRows( public Iterator> handleExpiredTimer( K key, TimerValues timerValues, ExpiredTimerInfo expiredTimerInfo) { List> outputs = new ArrayList<>(); - Instant watermark = new Instant(timerValues.getCurrentWatermarkInMs()); - SparkTimerInternals timerInternals = SparkTimerInternals.forWatermark(watermark); - restoreTimers(timerInternals); + SparkTimerInternals timerInternals = timerInternals(timerValues); + timerInternals.advanceWatermark(); DoFnRunner, 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); } diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTest.java index 26e900eba8b4..2ff0bfc2bf05 100644 --- a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTest.java +++ b/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 collected = StreamingTestUtils.collected(collectorId); + assertTrue(collected.contains("timer-output")); + } + private static final class KeySelectorFn extends DoFn> { @ProcessElement public void process(@Element String element, OutputReceiver> out) { @@ -239,4 +265,40 @@ public void onTimer(OutputReceiver out) { out.output("k1-timer"); } } + + private static final class TimerEmitterFn extends DoFn, KV> { + @TimerId("t") + private final TimerSpec timer = TimerSpecs.timer(TimeDomain.EVENT_TIME); + + @ProcessElement + public void process( + @Element KV element, + @Timestamp Instant ts, + @TimerId("t") Timer timer, + OutputReceiver> 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> out) { + out.outputWithTimestamp(KV.of("downstream", "timer-output"), ts); + } + } + + private static final class DownstreamFn extends DoFn, String> { + @StateId("state") + private final StateSpec> stateSpec = StateSpecs.value(); + + @ProcessElement + public void process( + @Element KV element, + @StateId("state") ValueState state, + OutputReceiver out) { + state.write(1); + out.output(element.getValue()); + } + } } diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/stateful/SparkTimerInternals.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/stateful/SparkTimerInternals.java index 4b904218b5a2..f234a3b865d8 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/stateful/SparkTimerInternals.java +++ b/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. */ From 086b65e485697922de70f6faee56f4e0ba2461fd Mon Sep 17 00:00:00 2001 From: Tobias Kaymak Date: Wed, 7 Oct 2026 11:38:43 +0000 Subject: [PATCH 2/3] Retrigger CI, AsyncWrapperTest flaky on Windows (#40448) From 9b647373a03006676134f2399a92c69cfaf5c709 Mon Sep 17 00:00:00 2001 From: Tobias Kaymak Date: Wed, 7 Oct 2026 13:13:25 +0000 Subject: [PATCH 3/3] Test SparkTimerInternals.forWatermarks --- .../spark/stateful/SparkTimerInternalsTest.java | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/runners/spark/src/test/java/org/apache/beam/runners/spark/stateful/SparkTimerInternalsTest.java b/runners/spark/src/test/java/org/apache/beam/runners/spark/stateful/SparkTimerInternalsTest.java index 9b65b683ff25..b8f8120dd675 100644 --- a/runners/spark/src/test/java/org/apache/beam/runners/spark/stateful/SparkTimerInternalsTest.java +++ b/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()); + } }