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. */ 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()); + } }