From 380d2fcb355a1040707b44fb44203855b7112ea8 Mon Sep 17 00:00:00 2001 From: Divyansh Yadav Date: Tue, 13 Oct 2026 14:11:01 +0530 Subject: [PATCH 1/2] [Spark 4] Clear the watermark of keys without timers or state in the streaming stateful ParDo --- .../state/BeamStatefulProcessor.java | 14 +++- .../streaming/StatefulParDoStreamingTest.java | 75 +++++++++++++++++++ 2 files changed, 87 insertions(+), 2 deletions(-) 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 2649f883af3f..cb3cae649514 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 @@ -120,13 +120,23 @@ public void init(OutputMode outputMode, TimeMode timeMode) { } private void storeWatermark(SparkTimerInternals timerInternals) { - checkStateNotNull(beamWatermark).update(timerInternals.currentInputWatermarkTime().getMillis()); + ValueState watermarkState = checkStateNotNull(beamWatermark); + // A key without timers and state, e.g. after all its windows expired, keeps nothing in the + // state store, otherwise an unbounded key space grows it forever. + if (timerInternals.getTimers().isEmpty() && !checkStateNotNull(beamState).exists()) { + if (watermarkState.exists()) { + watermarkState.clear(); + } + } else { + watermarkState.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. + // 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()) 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 2d3c127d72bc..de9e4921dab7 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 @@ -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,6 +73,9 @@ 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"; + + private static final int LATE_INDEX = 5; @After public void tearDown() { @@ -79,6 +83,7 @@ public void tearDown() { TestUnboundedSource.forget(TAG_RESTART); TestUnboundedSource.forget(TAG_CHAINED); TestUnboundedSource.forget(TAG_NULL_KEY); + TestUnboundedSource.forget(TAG_EXPIRY); } private static Pipeline pipeline( @@ -190,6 +195,44 @@ 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 ExpiryTestFn())) + .apply("Collect", ParDo.of(new StreamingTestUtils.CollectDoFn<>(collectorId))); + + assertEquals(PipelineResult.State.DONE, StreamingTestUtils.run(pipeline).getState()); + Set collected = StreamingTestUtils.collected(collectorId); + assertTrue(collected.contains("k1-0")); + // The late element finds the key cleared and is filtered like a first visit, see #40446. + assertTrue(collected.contains("k1-" + LATE_INDEX)); + 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> { @ProcessElement public void process(@Element String element, OutputReceiver> out) { @@ -311,6 +354,38 @@ public void onTimer(@Timestamp Instant ts, OutputReceiver> ou } } + private static final class ExpiryKeySelectorFn extends DoFn> { + @ProcessElement + public void process(@Element String element, OutputReceiver> out) { + int i = TestUnboundedSource.indexOf(element); + if (i == 0) { + out.output(KV.of("k1", element)); + } else if (i == LATE_INDEX) { + // Back into the first window of k1, expired by now. + out.outputWithTimestamp(KV.of("k1", element), new Instant(TestUnboundedSource.BASE_MILLIS)); + } + } + + @Override + public Duration getAllowedTimestampSkew() { + return Duration.millis(Long.MAX_VALUE); + } + } + + private static final class ExpiryTestFn extends DoFn, String> { + @StateId("seen") + private final StateSpec> seenSpec = StateSpecs.value(); + + @ProcessElement + public void process( + @Element KV element, + @StateId("seen") ValueState seenState, + OutputReceiver out) { + seenState.write(true); + out.output(element.getKey() + "-" + TestUnboundedSource.indexOf(element.getValue())); + } + } + private static final class DownstreamFn extends DoFn, String> { @StateId("state") private final StateSpec> stateSpec = StateSpecs.value(); From ce85c0626e2ce4ed638005d9576cd5658129e741 Mon Sep 17 00:00:00 2001 From: Divyansh Yadav Date: Fri, 16 Oct 2026 21:11:02 +0530 Subject: [PATCH 2/2] [Spark 4] Drop the late element from the expired key test and state the contract in storeWatermark --- .../state/BeamStatefulProcessor.java | 3 +- .../streaming/StatefulParDoStreamingTest.java | 30 ++----------------- 2 files changed, 3 insertions(+), 30 deletions(-) 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 cb3cae649514..7d794ef04506 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 @@ -121,8 +121,7 @@ public void init(OutputMode outputMode, TimeMode timeMode) { private void storeWatermark(SparkTimerInternals timerInternals) { ValueState watermarkState = checkStateNotNull(beamWatermark); - // A key without timers and state, e.g. after all its windows expired, keeps nothing in the - // state store, otherwise an unbounded key space grows it forever. + // 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(); 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 de9e4921dab7..18143ce4d924 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 @@ -75,8 +75,6 @@ public class StatefulParDoStreamingTest implements Serializable { private static final String TAG_NULL_KEY = "stateful-null-key"; private static final String TAG_EXPIRY = "stateful-expiry"; - private static final int LATE_INDEX = 5; - @After public void tearDown() { TestUnboundedSource.forget(TAG_BATCHES); @@ -208,14 +206,12 @@ public void testExpiredKeyLeavesNoState() throws Exception { .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 ExpiryTestFn())) + .apply("StatefulParDo", 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("k1-0")); - // The late element finds the key cleared and is filtered like a first visit, see #40446. - assertTrue(collected.contains("k1-" + LATE_INDEX)); + assertTrue(collected.contains(TestUnboundedSource.element(TAG_EXPIRY, 0, 0))); for (String stateVarName : Arrays.asList("beamState", "beamTimers", "beamWatermark")) { assertEquals(stateVarName, 0L, stateRows(checkpoint, stateVarName)); } @@ -360,30 +356,8 @@ public void process(@Element String element, OutputReceiver> int i = TestUnboundedSource.indexOf(element); if (i == 0) { out.output(KV.of("k1", element)); - } else if (i == LATE_INDEX) { - // Back into the first window of k1, expired by now. - out.outputWithTimestamp(KV.of("k1", element), new Instant(TestUnboundedSource.BASE_MILLIS)); } } - - @Override - public Duration getAllowedTimestampSkew() { - return Duration.millis(Long.MAX_VALUE); - } - } - - private static final class ExpiryTestFn extends DoFn, String> { - @StateId("seen") - private final StateSpec> seenSpec = StateSpecs.value(); - - @ProcessElement - public void process( - @Element KV element, - @StateId("seen") ValueState seenState, - OutputReceiver out) { - seenState.write(true); - out.output(element.getKey() + "-" + TestUnboundedSource.indexOf(element.getValue())); - } } private static final class DownstreamFn extends DoFn, String> {