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..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 @@ -120,13 +120,22 @@ 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 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 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..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 @@ -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,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 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 +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 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> { @ProcessElement public void process(@Element String element, OutputReceiver> out) { @@ -311,6 +350,16 @@ 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)); + } + } + } + private static final class DownstreamFn extends DoFn, String> { @StateId("state") private final StateSpec> stateSpec = StateSpecs.value();