From 989db7acbef508d294f7b931ad8e7f55b54062ed Mon Sep 17 00:00:00 2001 From: Tobias Kaymak Date: Sun, 25 Oct 2026 07:39:33 +0000 Subject: [PATCH] [#36841] Spark 4 streaming: detached reader for sibling scans of one source --- .../io/streaming/BeamReaderCache.java | 82 +++++++++++++++---- .../io/streaming/UnboundedSourceDataset.java | 54 ++++++++---- .../streaming/PAssertStreamingTest.java | 10 +++ 3 files changed, 110 insertions(+), 36 deletions(-) diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCache.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCache.java index 8348d04d289c..9ff5a4700f75 100644 --- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCache.java +++ b/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 CachedReader acquire( String key, long startEpoch, + long endEpoch, + @Nullable String runId, UnboundedSource source, PipelineOptions options, long idleTimeoutMillis, @@ -99,20 +105,27 @@ public static CachedReader 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 reused = (CachedReader) 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 CachedReader 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 CachedReader 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,30 +249,37 @@ public static final class CachedReader implements Closeable { private final UnboundedReader 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 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(); } @@ -262,6 +287,10 @@ public UnboundedReader 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; diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/UnboundedSourceDataset.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/UnboundedSourceDataset.java index 3ef35762748a..465ca4592afe 100644 --- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/UnboundedSourceDataset.java +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/UnboundedSourceDataset.java @@ -570,14 +570,16 @@ public PartitionReader createReader(InputPartition partition) { * *

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 the element type of the split */ static final class BeamPartitionReader implements PartitionReader { 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 split; @@ -615,6 +617,8 @@ static final class BeamPartitionReader implements PartitionReader in = pipeline.apply(Create.of("one", "two", "three")); + PCollection out = PCollectionList.of(in).and(in).apply(Flatten.pCollections()); + PAssert.that(out).containsInAnyOrder("one", "one", "two", "two", "three", "three"); + pipeline.run(options()); + } }