[Spark] Retain the last watermark of a source with no update in the batch

.github/trigger_files/beam_PostCommit_Java_PVR_Spark3_Streaming.json

@@ -1,5 +1,5 @@
 {
   "comment": "Modify this file in a trivial way to cause this test suite to run",
-  "modification": 8,
+  "modification": 9,
   "https://github.com/apache/beam/pull/39990": "removing dead code from FnApiDoFnRunner"
 }

.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": 3
+    "modification": 4
 }

CHANGES.md

@@ -93,6 +93,7 @@
 ## Bugfixes
 
 * (Java) Fixed the Spark runner firing processing-time timers in reverse timestamp order ([#39824](https://github.com/apache/beam/issues/39824)).
+* (Java) Fixed the Spark runner dropping the stored watermark of a streaming source with no update in a batch ([#39822](https://github.com/apache/beam/issues/39822)).
 * (Python) Fixed incorrect profiler options handling on portable runners ([#39613](https://github.com/apache/beam/issues/39613)).
 * (Java) KafkaIO dynamic reads no longer require the obsolete `beam_fn_api` experiment ([#29998](https://github.com/apache/beam/issues/29998)).
 * (Prism) Self-checkpointing splittable DoFns now resume after their requested delay instead of immediately, so polling SDFs no longer busy-spin ([#39848](https://github.com/apache/beam/issues/39848)).

runners/spark/src/main/java/org/apache/beam/runners/spark/stateful/SparkTimerInternals.java

@@ -17,8 +17,6 @@
  */
 package org.apache.beam.runners.spark.stateful;
 
-import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
-
 import java.util.Collection;
 import java.util.Collections;
 import java.util.Comparator;
@@ -83,14 +81,13 @@ public static SparkTimerInternals forStreamFromSources(
             slowestHighWatermark.isBefore(sparkWatermarks.getHighWatermark())
                 ? slowestHighWatermark
                 : sparkWatermarks.getHighWatermark();
-        if (synchronizedProcessingTime == null) {
-          // firstime set.
+        // an idle source keeps its last reported time, so sources can disagree here; the
+        // synchronized processing time is the one every source has reached.
+        if (synchronizedProcessingTime == null
+            || sparkWatermarks
+                .getSynchronizedProcessingTime()
+                .isBefore(synchronizedProcessingTime)) {
           synchronizedProcessingTime = sparkWatermarks.getSynchronizedProcessingTime();
-        } else {
-          // assert on following.
-          checkArgument(
-              sparkWatermarks.getSynchronizedProcessingTime().equals(synchronizedProcessingTime),
-              "Synchronized time is expected to keep synchronized across sources.");
         }
       }
     }

runners/spark/src/main/java/org/apache/beam/runners/spark/util/GlobalWatermarkHolder.java

@@ -175,22 +175,36 @@ public static void advance() {
   /**
    * Computes the next watermark values per source id.
    *
-   * @return The new watermarks values or null if no source has reported its progress.
+   * <p>Sources with no queued update keep their last known watermarks.
+   *
+   * @return The new watermarks values, or an empty map if no source has reported its progress.
    */
   private static Map<Integer, SparkWatermarks> computeNewWatermarks(BlockManager blockManager) {
 
     if (sourceTimes.isEmpty()) {
       return new HashMap<>();
     }
 
-    // update all sources' watermarks into the new broadcast.
+    final Map<Integer, SparkWatermarks> currentWatermarks = initWatermarks(blockManager);
+
+    // Update the sources that reported progress over a copy of the stored map, so a source
+    // with no update in this batch keeps its last known watermarks. A source whose high
+    // watermark already reached the end of time has finished, so it is dropped before its
+    // last low watermark pins the minimum.
     final Map<Integer, SparkWatermarks> newValues = new HashMap<>();
+    for (Map.Entry<Integer, SparkWatermarks> current : currentWatermarks.entrySet()) {
+      if (current.getValue().getHighWatermark().isBefore(BoundedWindow.TIMESTAMP_MAX_VALUE)) {
+        newValues.put(current.getKey(), current.getValue());
+      }
+    }
+    boolean hasUpdates = false;
 
     for (final Map.Entry<Integer, Queue<SparkWatermarks>> watermarkInfo : sourceTimes.entrySet()) {
 
       if (watermarkInfo.getValue().isEmpty()) {
         continue;
       }
+      hasUpdates = true;
 
       final Integer sourceId = watermarkInfo.getKey();
 
@@ -199,8 +213,6 @@ private static Map<Integer, SparkWatermarks> computeNewWatermarks(BlockManager b
       Instant currentHighWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE;
       Instant currentSynchronizedProcessingTime = BoundedWindow.TIMESTAMP_MIN_VALUE;
 
-      final Map<Integer, SparkWatermarks> currentWatermarks = initWatermarks(blockManager);
-
       if (currentWatermarks.containsKey(sourceId)) {
         final SparkWatermarks currentTimes = currentWatermarks.get(sourceId);
         currentLowWatermark = currentTimes.getLowWatermark();
@@ -240,7 +252,7 @@ private static Map<Integer, SparkWatermarks> computeNewWatermarks(BlockManager b
           new SparkWatermarks(nextLowWatermark, nextHighWatermark, nextSynchronizedProcessingTime));
     }
 
-    return newValues;
+    return hasUpdates ? newValues : new HashMap<>();
   }
 
   private static void writeRemoteWatermarkBlock(

runners/spark/src/test/java/org/apache/beam/runners/spark/GlobalWatermarkHolderTest.java

@@ -17,12 +17,14 @@
  */
 package org.apache.beam.runners.spark;
 
+import static org.hamcrest.CoreMatchers.notNullValue;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.core.IsEqual.equalTo;
 
 import org.apache.beam.runners.spark.util.GlobalWatermarkHolder;
 import org.apache.beam.runners.spark.util.GlobalWatermarkHolder.SparkWatermarks;
 import org.apache.beam.sdk.testing.RegexMatcher;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
 import org.joda.time.Duration;
 import org.joda.time.Instant;
 import org.junit.ClassRule;
@@ -134,4 +136,74 @@ public void testMultiSource() {
     assertThat(watermarksForSource2.getLowWatermark(), equalTo(instant.plus(Duration.millis(3))));
     assertThat(watermarksForSource2.getHighWatermark(), equalTo(instant.plus(Duration.millis(6))));
   }
+
+  @Test
+  public void testWatermarkRetainedForSourceWithoutUpdate() {
+    Instant instant = new Instant(0);
+
+    GlobalWatermarkHolder.add(
+        1,
+        new SparkWatermarks(
+            instant.plus(Duration.millis(5)), instant.plus(Duration.millis(10)), instant));
+    GlobalWatermarkHolder.add(
+        2,
+        new SparkWatermarks(
+            instant.plus(Duration.millis(3)), instant.plus(Duration.millis(6)), instant));
+
+    GlobalWatermarkHolder.advance();
+
+    // only source 1 reports progress in the next batch.
+    GlobalWatermarkHolder.add(
+        1,
+        new SparkWatermarks(
+            instant.plus(Duration.millis(8)),
+            instant.plus(Duration.millis(12)),
+            instant.plus(Duration.millis(100))));
+
+    GlobalWatermarkHolder.advance();
+
+    // source 1 advanced.
+    SparkWatermarks watermarksForSource1 = GlobalWatermarkHolder.get(0L).get(1);
+    assertThat(watermarksForSource1.getLowWatermark(), equalTo(instant.plus(Duration.millis(8))));
+    assertThat(watermarksForSource1.getHighWatermark(), equalTo(instant.plus(Duration.millis(12))));
+
+    // source 2 had no update; its last known watermarks must be retained.
+    SparkWatermarks watermarksForSource2 = GlobalWatermarkHolder.get(0L).get(2);
+    assertThat(watermarksForSource2, notNullValue());
+    assertThat(watermarksForSource2.getLowWatermark(), equalTo(instant.plus(Duration.millis(3))));
+    assertThat(watermarksForSource2.getHighWatermark(), equalTo(instant.plus(Duration.millis(6))));
+  }
+
+  @Test
+  public void testCompletedSourceAgesOut() {
+    Instant instant = new Instant(0);
+
+    // source 1 completes: its high watermark reaches the end of time.
+    GlobalWatermarkHolder.add(
+        1,
+        new SparkWatermarks(
+            instant.plus(Duration.millis(5)), BoundedWindow.TIMESTAMP_MAX_VALUE, instant));
+    GlobalWatermarkHolder.add(
+        2,
+        new SparkWatermarks(
+            instant.plus(Duration.millis(3)), instant.plus(Duration.millis(6)), instant));
+
+    GlobalWatermarkHolder.advance();
+
+    // the batch that completed the source still reports it.
+    assertThat(GlobalWatermarkHolder.get(0L).get(1), notNullValue());
+
+    GlobalWatermarkHolder.add(
+        2,
+        new SparkWatermarks(
+            instant.plus(Duration.millis(8)),
+            instant.plus(Duration.millis(12)),
+            instant.plus(Duration.millis(100))));
+
+    GlobalWatermarkHolder.advance();
+
+    // a completed source is dropped so its last low watermark stops pinning the minimum.
+    assertThat(GlobalWatermarkHolder.get(0L).containsKey(1), equalTo(false));
+    assertThat(GlobalWatermarkHolder.get(0L).get(2), notNullValue());
+  }
 }

runners/spark/src/test/java/org/apache/beam/runners/spark/translation/streaming/CreateStreamTest.java

@@ -310,7 +310,9 @@ public void testFlattenedWithWatermarkHold() throws IOException {
                 TimestampedValue.of(1, instant),
                 TimestampedValue.of(2, instant),
                 TimestampedValue.of(3, instant))
-            .advanceWatermarkForNextBatch(instant.plus(Duration.standardMinutes(10)));
+            .advanceWatermarkForNextBatch(instant.plus(Duration.standardMinutes(10)))
+            .emptyBatch()
+            .advanceNextBatchWatermarkToInfinity();
     CreateStream<Integer> source2 =
         CreateStream.of(VarIntCoder.of(), batchDuration())
             .emptyBatch()

runners/spark/src/test/java/org/apache/beam/runners/spark/translation/streaming/StatefulStreamingParDoEvaluatorTest.java

@@ -278,8 +278,10 @@ public void shouldTriggerProcessingTimeTimerWithSparseKey() {
         CreateStream.of(coder, batchDuration(p)).advanceWatermarkForNextBatch(instant);
 
     sparseStream =
-        sparseStream.nextBatch(
-            TimestampedValue.of(KV.of(sparseKey, 0), instant.plus(Duration.millis(1000L))));
+        sparseStream
+            .nextBatch(
+                TimestampedValue.of(KV.of(sparseKey, 0), instant.plus(Duration.millis(1000L))))
+            .advanceNextBatchWatermarkToInfinity();
 
     final PCollection<KV<Integer, Integer>> sparsePCollection =
         p.apply("Create Sparse Key Stream", sparseStream);