This is an automated email from the ASF dual-hosted git repository.
Abacn pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 0724ef9ce58 [Spark 4] Clear the watermark of keys without timers or
state in the streaming stateful ParDo (#40550)
0724ef9ce58 is described below
commit 0724ef9ce58b86edb4b83a89dc4825f319edd7b7
Author: Divyansh Yadav <[email protected]>
AuthorDate: Wed Oct 21 04:20:37 2026 +0530
[Spark 4] Clear the watermark of keys without timers or state in the
streaming stateful ParDo (#40550)
---
.../streaming/state/BeamStatefulProcessor.java | 13 +++++-
.../streaming/StatefulParDoStreamingTest.java | 49 ++++++++++++++++++++++
2 files changed, 60 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 2649f883af3..7d794ef0450 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,22 @@ public class BeamStatefulProcessor<K, V, OutputT>
}
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())
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 2d3c127d72b..18143ce4d92 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.assertEquals;
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,7 @@ 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() {
@@ -79,6 +81,7 @@ public class StatefulParDoStreamingTest implements
Serializable {
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 class StatefulParDoStreamingTest implements
Serializable {
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 class StatefulParDoStreamingTest implements
Serializable {
}
}
+ 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();