[#36841] Spark 4 streaming: detached reader for sibling scans of one source

runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCache.java

@@ -81,14 +81,20 @@ interface MarkRestorer {
 
   /**
    * Returns the reader for {@code key} positioned at {@code startEpoch}, reusing the cached one if
-   * it is there, restoring from the durable mark otherwise. A zero length durable mark means a
-   * fresh start. {@code committedEpoch} supplies the epoch Spark last committed, -1 if unknown.
+   * it is there, restoring from the durable mark otherwise. A sibling scan of the same batch in the
+   * same query run {@code runId} gets a detached reader restored from {@code startEpoch}. Sibling
+   * scans read independently, their batches agree only when the source yields the same records up
+   * to the same cut, a deadline cut can differ per sibling, see #40786. A zero length durable mark
+   * means a fresh start. {@code committedEpoch} supplies the epoch Spark last committed, -1 if
+   * unknown.
    *
    * @throws IllegalStateException if {@code startEpoch > 0} and no durable mark exists
    */
   public static <T> CachedReader<T> acquire(
       String key,
       long startEpoch,
+      long endEpoch,
+      @Nullable String runId,
       UnboundedSource<T, ?> source,
       PipelineOptions options,
       long idleTimeoutMillis,
@@ -99,20 +105,27 @@ public static <T> CachedReader<T> acquire(
     closeIdle();
     synchronized (lock(key)) {
       CachedReader<?> existing = READERS.get(key);
+      boolean detached = false;
+      boolean sentinelEmitted = false;
       if (existing != null) {
-        if (existing.beginBatch(startEpoch)) {
+        if (existing.beginBatch(startEpoch, runId)) {
           existing.finalizePendingMark(key);
           @SuppressWarnings("unchecked") // one source per key, its element type never changes
           CachedReader<T> reused = (CachedReader<T>) existing;
           return reused;
         }
-        LOG.info(
-            "Cached Beam reader {} is at epoch {}, batch starts at {}, restoring from the durable"
-                + " mark.",
-            key,
-            existing.positionEpoch(),
-            startEpoch);
-        invalidate(key);
+        if (existing.isSibling(startEpoch, endEpoch, runId)) {
+          detached = true;
+          sentinelEmitted = existing.sentinelDone();
+        } else {
+          LOG.info(
+              "Cached Beam reader {} is at epoch {}, batch starts at {}, restoring from the durable"
+                  + " mark.",
+              key,
+              existing.positionEpoch(),
+              startEpoch);
+          invalidate(key);
+        }
       }
       byte[] codedMark = restorer.restore();
       if (codedMark == null && startEpoch > 0) {
@@ -123,7 +136,8 @@ public static <T> CachedReader<T> acquire(
         codedMark = null;
       }
       LOG.info(
-          "Creating Beam reader {} at epoch {} ({} mark).",
+          "Creating {}Beam reader {} at epoch {} ({} mark).",
+          detached ? "detached " : "",
           key,
           startEpoch,
           codedMark == null ? "no" : "restored");
@@ -133,9 +147,13 @@ public static <T> CachedReader<T> acquire(
               startEpoch,
               codedMark,
               idleTimeoutMillis,
-              committedEpoch);
-      created.beginBatch(startEpoch);
-      READERS.put(key, created);
+              committedEpoch,
+              detached,
+              sentinelEmitted);
+      created.beginBatch(startEpoch, runId);
+      if (!detached) {
+        READERS.put(key, created);
+      }
       return created;
     }
   }
@@ -231,37 +249,48 @@ public static final class CachedReader<T> implements Closeable {
     private final UnboundedReader<T> reader;
     private final long idleTimeoutMillis;
     private final LongSupplier committedEpoch;
+    private final boolean detached;
     private boolean started;
     private boolean exhausted;
     private boolean positionExhausted;
     private boolean doneEligible;
     private boolean doneWritten;
     private boolean sentinelEmitted;
+    private boolean sentinelDone;
     private boolean inBatch;
     private boolean moved;
     private long positionEpoch;
     private byte @Nullable [] positionMark;
     private @Nullable CheckpointMark pendingMark;
     private long lastUsedMillis;
+    private @Nullable String runId;
 
     CachedReader(
         UnboundedReader<T> reader,
         long positionEpoch,
         byte @Nullable [] positionMark,
         long idleTimeoutMillis,
-        LongSupplier committedEpoch) {
+        LongSupplier committedEpoch,
+        boolean detached,
+        boolean sentinelEmitted) {
       this.reader = reader;
       this.positionEpoch = positionEpoch;
       this.positionMark = positionMark;
       this.idleTimeoutMillis = idleTimeoutMillis;
       this.committedEpoch = committedEpoch;
+      this.detached = detached;
+      this.sentinelEmitted = sentinelEmitted;
       this.lastUsedMillis = System.currentTimeMillis();
     }
 
     public UnboundedReader<T> reader() {
       return reader;
     }
 
+    boolean isDetached() {
+      return detached;
+    }
+
     public synchronized boolean startOrAdvance() throws IOException {
       moved = true;
       boolean hasElement = started ? reader.advance() : reader.start();
@@ -284,6 +313,10 @@ public synchronized boolean claimSentinel(boolean emitSentinel) {
       return true;
     }
 
+    synchronized boolean sentinelDone() {
+      return sentinelDone;
+    }
+
     /**
      * Whether a committed batch exhausted this reader and its done marker has not been written yet.
      */
@@ -308,18 +341,31 @@ public synchronized long positionEpoch() {
     }
 
     /**
-     * Claims the reader for a batch starting at {@code epoch}, false if it cannot continue there.
+     * Claims the reader for a batch starting at {@code epoch} in query run {@code runId}, false if
+     * it cannot continue there.
      */
-    synchronized boolean beginBatch(long epoch) {
-      if (positionEpoch != epoch || moved) {
+    synchronized boolean beginBatch(long epoch, @Nullable String runId) {
+      if (inBatch || positionEpoch != epoch || moved) {
         return false;
       }
       doneEligible |= positionExhausted;
+      sentinelDone |= sentinelEmitted;
       inBatch = true;
+      this.runId = runId;
       lastUsedMillis = System.currentTimeMillis();
       return true;
     }
 
+    /**
+     * Whether this entry is in or just finished the batch from {@code startEpoch} to {@code
+     * endEpoch} of query run {@code runId}.
+     */
+    synchronized boolean isSibling(long startEpoch, long endEpoch, @Nullable String runId) {
+      return runId != null
+          && runId.equals(this.runId)
+          && ((inBatch && positionEpoch == startEpoch) || (!inBatch && positionEpoch == endEpoch));
+    }
+
     /** Records a completed batch, the reader is positioned at {@code endEpoch} from now on. */
     synchronized void endBatch(long endEpoch, @Nullable CheckpointMark mark, byte[] codedMark) {
       positionEpoch = endEpoch;

runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/UnboundedSourceDataset.java

@@ -570,14 +570,16 @@ public PartitionReader<InternalRow> createReader(InputPartition partition) {
    *
    * <p>The batch ends at the record quota or at the deadline. The reader then writes its checkpoint
    * mark durably at the end epoch and stays in {@link BeamReaderCache} for the next batch. A failed
-   * mark write fails the task. An attempt Spark killed or failed writes nothing and its reader is
-   * dropped, the retry restores from the durable mark at the start epoch.
+   * mark write fails the task and drops the reader. An attempt Spark killed or failed writes
+   * nothing and its reader is dropped, the retry restores from the durable mark at the start epoch.
    *
    * @param <T> the element type of the split
    */
   static final class BeamPartitionReader<T> implements PartitionReader<InternalRow> {
 
     private static final Duration INITIAL_BACKOFF = Duration.millis(10);
+    // StreamExecution sets the job group to the query run id.
+    private static final String JOB_GROUP_ID = "spark.jobGroup.id";
 
     private final String key;
     private final UnboundedSource<T, ?> split;
@@ -615,6 +617,8 @@ static final class BeamPartitionReader<T> implements PartitionReader<InternalRow
           BeamReaderCache.acquire(
               key,
               startEpoch,
+              endEpoch,
+              runId(),
               split,
               options,
               partition.readerIdleTimeoutMillis,
@@ -681,35 +685,49 @@ private boolean endOfBatch(boolean discarded) throws IOException {
     }
 
     /**
-     * Ends the batch once. A discarded attempt drops the reader and writes nothing. A reader that
-     * was never started has not moved, its start mark is written forward, an empty file standing
-     * for a fresh start.
+     * Ends the batch once. A detached reader closes without writing marks. A discarded attempt
+     * drops the reader and writes nothing. A reader that was never started has not moved, its start
+     * mark is written forward, an empty file standing for a fresh start.
      */
     private void endBatch(boolean discarded) throws IOException {
       if (batchEnded) {
         return;
       }
       batchEnded = true;
+      if (cached.isDetached()) {
+        cached.close();
+        return;
+      }
       if (discarded) {
         LOG.info("Attempt for Beam reader {} was discarded, dropping the reader.", key);
         BeamReaderCache.invalidate(key);
         return;
       }
-      if (!cached.started()) {
-        byte[] startMark = cached.positionMark();
-        byte[] codedMark = startMark == null ? new byte[0] : startMark;
+      try {
+        if (!cached.started()) {
+          byte[] startMark = cached.positionMark();
+          byte[] codedMark = startMark == null ? new byte[0] : startMark;
+          checkpoint.writeMark(splitId, endEpoch, codedMark);
+          cached.endBatch(endEpoch, null, codedMark);
+          return;
+        }
+        CheckpointMark mark = cached.reader().getCheckpointMark();
+        byte[] codedMark = encodeMark(split, mark);
         checkpoint.writeMark(splitId, endEpoch, codedMark);
-        cached.endBatch(endEpoch, null, codedMark);
-        return;
-      }
-      CheckpointMark mark = cached.reader().getCheckpointMark();
-      byte[] codedMark = encodeMark(split, mark);
-      checkpoint.writeMark(splitId, endEpoch, codedMark);
-      if (cached.shouldWriteDone()) {
-        checkpoint.writeDone(splitId);
+        if (cached.shouldWriteDone()) {
+          checkpoint.writeDone(splitId);
+        }
+        cached.endBatch(endEpoch, mark, codedMark);
+        LOG.debug("Beam reader {} read {} record(s) up to epoch {}.", key, recordsRead, endEpoch);
+      } catch (IOException | RuntimeException e) {
+        BeamReaderCache.invalidate(key);
+        throw e;
       }
-      cached.endBatch(endEpoch, mark, codedMark);
-      LOG.debug("Beam reader {} read {} record(s) up to epoch {}.", key, recordsRead, endEpoch);
+    }
+
+    private static @Nullable String runId() {
+      TaskContext context = TaskContext.get();
+      return context == null ? null : context.getLocalProperty(JOB_GROUP_ID);
     }
 
     private static boolean attemptDiscarded() {

runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/PAssertStreamingTest.java

@@ -31,9 +31,11 @@
 import org.apache.beam.sdk.testing.PAssert;
 import org.apache.beam.sdk.testing.TestPipeline;
 import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.transforms.Flatten;
 import org.apache.beam.sdk.transforms.windowing.FixedWindows;
 import org.apache.beam.sdk.transforms.windowing.Window;
 import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionList;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Throwables;
 import org.joda.time.Duration;
 import org.junit.After;
@@ -146,4 +148,12 @@ public void testCreateMultipleElements() throws Exception {
     PAssert.that(out).containsInAnyOrder("one", "two", "three");
     pipeline.run(options());
   }
+
+  @Test
+  public void testFlattenSameSource() throws Exception {
+    PCollection<String> in = pipeline.apply(Create.of("one", "two", "three"));
+    PCollection<String> out = PCollectionList.of(in).and(in).apply(Flatten.pCollections());
+    PAssert.that(out).containsInAnyOrder("one", "one", "two", "two", "three", "three");
+    pipeline.run(options());
+  }
 }