From a218d94f0a8d1fb029d10358b8be1c6e3b6bbdce Mon Sep 17 00:00:00 2001 From: Danny McCormick Date: Fri, 13 Nov 2026 07:52:41 +0000 Subject: [PATCH 1/2] Synchronize SingleEmitInputDStream batches with GlobalWatermarkHolder --- .../runners/spark/SparkPipelineRunner.java | 4 ++++ .../translation/SingleEmitInputDStream.java | 18 ++++++++++++++++++ .../spark/util/GlobalWatermarkHolder.java | 11 ++++++++--- 3 files changed, 30 insertions(+), 3 deletions(-) diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkPipelineRunner.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkPipelineRunner.java index 4e4841da13eb..a31c2207be45 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkPipelineRunner.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkPipelineRunner.java @@ -123,8 +123,12 @@ public PortablePipelineResult run(RunnerApi.Pipeline pipeline, JobInfo jobInfo) // Initialize accumulators. Only the DStream streaming path uses the metrics checkpoint. MetricsEnvironment.setMetricsSupported(true); + MetricsAccumulator.clear(); MetricsAccumulator.init( pipelineOptions, jsc, !useStructuredStreaming && pipelineOptions.isStreaming()); + if (useDStreams) { + GlobalWatermarkHolder.clear(); + } final SparkTranslationContext context = translator.createTranslationContext(jsc, pipelineOptions, jobInfo); diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SingleEmitInputDStream.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SingleEmitInputDStream.java index 3cde41d2812e..3a18f16e14d7 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SingleEmitInputDStream.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SingleEmitInputDStream.java @@ -17,6 +17,7 @@ */ package org.apache.beam.runners.spark.translation; +import org.apache.beam.runners.spark.util.GlobalWatermarkHolder; import org.apache.spark.api.java.JavaSparkContext$; import org.apache.spark.rdd.RDD; import org.apache.spark.streaming.StreamingContext; @@ -42,6 +43,7 @@ public class SingleEmitInputDStream extends ConstantInputDStream { private boolean emitted = false; + private transient long lastValidTimeMs = 0; public SingleEmitInputDStream(StreamingContext ssc, RDD rdd) { super(ssc, rdd, JavaSparkContext$.MODULE$.fakeClassTag()); @@ -49,6 +51,9 @@ public SingleEmitInputDStream(StreamingContext ssc, RDD rdd) { @Override public Option> compute(Time validTime) { + if (validTime != null) { + waitForLastBatch(validTime); + } if (this.emitted) { return Option.apply(this.emptyRDD()); } else { @@ -57,6 +62,19 @@ public Option> compute(Time validTime) { } } + private void waitForLastBatch(Time validTime) { + while (GlobalWatermarkHolder.hasWatermarks() + && GlobalWatermarkHolder.getLastWatermarkedBatchTime() < lastValidTimeMs) { + try { + Thread.sleep(10); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + break; + } + } + lastValidTimeMs = validTime.milliseconds(); + } + private RDD emptyRDD() { return this.context().sparkContext().emptyRDD(JavaSparkContext$.MODULE$.fakeClassTag()); } diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/util/GlobalWatermarkHolder.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/util/GlobalWatermarkHolder.java index 3f85247644a7..842f36184d80 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/util/GlobalWatermarkHolder.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/util/GlobalWatermarkHolder.java @@ -75,7 +75,7 @@ public class GlobalWatermarkHolder { private static volatile LoadingCache> watermarkCache = null; private static volatile long lastWatermarkedBatchTime = 0; - public static void add(int sourceId, SparkWatermarks sparkWatermarks) { + public static synchronized void add(int sourceId, SparkWatermarks sparkWatermarks) { Queue timesQueue = sourceTimes.get(sourceId); if (timesQueue == null) { timesQueue = new ConcurrentLinkedQueue<>(); @@ -85,7 +85,7 @@ public static void add(int sourceId, SparkWatermarks sparkWatermarks) { } @VisibleForTesting - public static void addAll(Map> sourceTimes) { + public static synchronized void addAll(Map> sourceTimes) { for (Map.Entry> en : sourceTimes.entrySet()) { int sourceId = en.getKey(); Queue timesQueue = en.getValue(); @@ -95,6 +95,10 @@ public static void addAll(Map> sourceTimes) { } } + public static synchronized boolean hasWatermarks() { + return !sourceTimes.isEmpty() || driverNodeWatermarks != null; + } + public static long getLastWatermarkedBatchTime() { return lastWatermarkedBatchTime; } @@ -102,7 +106,7 @@ public static long getLastWatermarkedBatchTime() { /** * Returns the {@link Broadcast} containing the {@link SparkWatermarks} mapped to their sources. */ - public static Map get(Long cacheInterval) { + public static synchronized Map get(Long cacheInterval) { if (canBypassRemoteWatermarkFetching()) { /* driverNodeWatermarks != null => @@ -313,6 +317,7 @@ public static synchronized void clear() { sourceTimes.clear(); lastWatermarkedBatchTime = 0; writeLocalWatermarkCopy(null); + watermarkCache = null; final SparkEnv sparkEnv = SparkEnv.get(); if (sparkEnv != null) { final BlockManager blockManager = sparkEnv.blockManager(); From 03c2a65c76d140a8c009d9edcfde0f8d2f15aa10 Mon Sep 17 00:00:00 2001 From: Danny McCormick Date: Fri, 13 Nov 2026 11:42:00 +0000 Subject: [PATCH 2/2] Enable SingleEmitInputDStream watermark sync only for portable streaming and clear GlobalWatermarkHolder after streaming runs --- .../runners/spark/SparkPipelineRunner.java | 28 +++++++++++-------- .../beam/runners/spark/TestSparkRunner.java | 1 + .../translation/SingleEmitInputDStream.java | 8 +++++- ...rkStreamingPortablePipelineTranslator.java | 9 ++++-- 4 files changed, 31 insertions(+), 15 deletions(-) diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkPipelineRunner.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkPipelineRunner.java index a31c2207be45..dcbceaf9ea80 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkPipelineRunner.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkPipelineRunner.java @@ -176,19 +176,25 @@ public PortablePipelineResult run(RunnerApi.Pipeline pipeline, JobInfo jobInfo) final Future submissionFuture = executorService.submit( () -> { - translator.translate(fusedPipeline, context); - LOG.info( - "Job {}: Pipeline translated successfully. Computing outputs", jobInfo.jobId()); - context.computeOutputs(); - - jssc.start(); + GlobalWatermarkHolder.clear(); try { - jssc.awaitTerminationOrTimeout(timeout); - } catch (InterruptedException e) { - LOG.warn("Streaming context interrupted, shutting down.", e); + translator.translate(fusedPipeline, context); + LOG.info( + "Job {}: Pipeline translated successfully. Computing outputs", + jobInfo.jobId()); + context.computeOutputs(); + + jssc.start(); + try { + jssc.awaitTerminationOrTimeout(timeout); + } catch (InterruptedException e) { + LOG.warn("Streaming context interrupted, shutting down.", e); + } + jssc.stop(); + LOG.info("Job {} finished.", jobInfo.jobId()); + } finally { + GlobalWatermarkHolder.clear(); } - jssc.stop(); - LOG.info("Job {} finished.", jobInfo.jobId()); }); result = new SparkPipelineResult.PortableStreamingMode(submissionFuture, jssc); } else { diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/TestSparkRunner.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/TestSparkRunner.java index 22e25e5272a2..b951ea9ce27b 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/TestSparkRunner.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/TestSparkRunner.java @@ -105,6 +105,7 @@ public SparkPipelineResult run(Pipeline pipeline) { finishState, isOneOf(PipelineResult.State.STOPPED, PipelineResult.State.DONE)); } finally { + GlobalWatermarkHolder.clear(); try { if (testSparkOptions.isDeleteCheckpointDir()) { // cleanup checkpoint dir. diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SingleEmitInputDStream.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SingleEmitInputDStream.java index 3a18f16e14d7..7f72e8a777c1 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SingleEmitInputDStream.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SingleEmitInputDStream.java @@ -43,15 +43,21 @@ public class SingleEmitInputDStream extends ConstantInputDStream { private boolean emitted = false; + private final boolean forceWatermarkSync; private transient long lastValidTimeMs = 0; public SingleEmitInputDStream(StreamingContext ssc, RDD rdd) { + this(ssc, rdd, false); + } + + public SingleEmitInputDStream(StreamingContext ssc, RDD rdd, boolean forceWatermarkSync) { super(ssc, rdd, JavaSparkContext$.MODULE$.fakeClassTag()); + this.forceWatermarkSync = forceWatermarkSync; } @Override public Option> compute(Time validTime) { - if (validTime != null) { + if (forceWatermarkSync && validTime != null && validTime.milliseconds() > lastValidTimeMs) { waitForLastBatch(validTime); } if (this.emitted) { diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkStreamingPortablePipelineTranslator.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkStreamingPortablePipelineTranslator.java index 7bfad22d9a52..2657ae218189 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkStreamingPortablePipelineTranslator.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkStreamingPortablePipelineTranslator.java @@ -158,7 +158,8 @@ private static void translateImpulse( .map(CoderHelpers.fromByteFunction(windowCoder)); final SingleEmitInputDStream> inputDStream = - new SingleEmitInputDStream<>(context.getStreamingContext().ssc(), emptyByteArrayRDD.rdd()); + new SingleEmitInputDStream<>( + context.getStreamingContext().ssc(), emptyByteArrayRDD.rdd(), true); final JavaDStream> stream = JavaDStream.fromDStream(inputDStream, JavaSparkContext$.MODULE$.fakeClassTag()); @@ -308,7 +309,7 @@ private static void translateFlatten( if (inputsMap.isEmpty()) { final JavaRDD> emptyRDD = context.getSparkContext().emptyRDD(); final SingleEmitInputDStream> singleEmitInputDStream = - new SingleEmitInputDStream<>(context.getStreamingContext().ssc(), emptyRDD.rdd()); + new SingleEmitInputDStream<>(context.getStreamingContext().ssc(), emptyRDD.rdd(), true); unifiedStreams = JavaDStream.fromDStream(singleEmitInputDStream, JavaSparkContext$.MODULE$.fakeClassTag()); } else { @@ -323,7 +324,9 @@ private static void translateFlatten( // create a single RDD stream. final SingleEmitInputDStream> singleEmitInputDStream = new SingleEmitInputDStream>( - context.getStreamingContext().ssc(), ((BoundedDataset) dataset).getRDD().rdd()); + context.getStreamingContext().ssc(), + ((BoundedDataset) dataset).getRDD().rdd(), + true); final JavaDStream> dStream = JavaDStream.fromDStream( singleEmitInputDStream, JavaSparkContext$.MODULE$.fakeClassTag());