[Spark] Support SDF self-checkpointing and bundle finalization on the portable runner

Eliaaazzz · Eliaaazzz · commit 17df8a30cfd4 · 2026-07-15T15:04:30.000+10:00
The portable Spark runner never registered a bundle checkpoint handler, so a
splittable DoFn that self-checkpoints failed on its first bundle and unbounded
reads could not run at all.

In batch, the executable stage function now collects residual roots and re-feeds
them in fresh bundles until the SDK returns none, resuming each residual at its
requested time. Unbounded residuals are rejected, since draining them would never
terminate. The stage also registers an in-memory bundle finalizer.

In streaming, residuals leave the stage on a reserved union tag and are relayed
on the driver: they are held until their resume time, fed back as stage input in
a later micro-batch, and their output watermarks advance the stage's watermark so
downstream event-time windows fire. The relay bounds its watermark by the stage's
upstream sources, and is keyed per job so concurrent jobs stay isolated.

Unskips the splittable DoFn and bundle finalization tests for the Spark runner.
diff --git a/.github/trigger_files/beam_PostCommit_Java_PVR_Spark3_Streaming.json b/.github/trigger_files/beam_PostCommit_Java_PVR_Spark3_Streaming.json
@@ -1,4 +1,4 @@
 {
   "comment": "Modify this file in a trivial way to cause this test suite to run",
-  "modification": 6
+  "modification": 7
 }
diff --git a/.github/trigger_files/beam_PostCommit_Java_PVR_Spark_Batch.json b/.github/trigger_files/beam_PostCommit_Java_PVR_Spark_Batch.json
@@ -1,4 +1,4 @@
 {
   "comment": "Modify this file in a trivial way to cause this test suite to run",
-  "modification": 6
+  "modification": 7
 }
diff --git 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": 1
+    "modification": 2
 }
diff --git a/.github/trigger_files/beam_PostCommit_Python_ValidatesRunner_Spark.json b/.github/trigger_files/beam_PostCommit_Python_ValidatesRunner_Spark.json
@@ -2,5 +2,6 @@
   "https://github.com/apache/beam/pull/34830": "testing",
   "https://github.com/apache/beam/issues/35429": "testing",
   "trigger-2026-04-04": "portable_runner expand_sdf opt-in",
-  "https://github.com/apache/beam/pull/38892": "UnboundedSource portable VR test"
+  "https://github.com/apache/beam/pull/38892": "UnboundedSource portable VR test",
+  "https://github.com/apache/beam/issues/19468": "SDF self-checkpointing and bundle finalization"
 }
diff --git a/CHANGES.md b/CHANGES.md
@@ -69,6 +69,8 @@
 ## New Features / Improvements
 
 * (Python) Removed the `envoy-data-plane` (and transitive `betterproto`) dependency; `EnvoyRateLimiter` now uses a small vendored protobuf definition instead, resolving dependency conflicts for downstream projects ([#37854](https://github.com/apache/beam/issues/37854)).
+* Splittable DoFn self-checkpointing is now supported on the portable Spark runner, including unbounded reads in streaming mode ([#19468](https://github.com/apache/beam/issues/19468)).
+* Bundle finalization is now supported on the portable Spark runner ([#19517](https://github.com/apache/beam/issues/19517)).
 * X feature added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)).
 
 ## Breaking Changes
diff --git a/runners/spark/job-server/spark_job_server.gradle b/runners/spark/job-server/spark_job_server.gradle
@@ -199,11 +199,9 @@ def portableValidatesRunnerTask(String name, boolean streaming, boolean docker,
         excludeCategories 'org.apache.beam.sdk.testing.UsesKeyInParDo'
         excludeCategories 'org.apache.beam.sdk.testing.UsesOnWindowExpiration'
         excludeCategories 'org.apache.beam.sdk.testing.UsesTestStream'
-        // TODO (https://github.com/apache/beam/issues/19468) SplittableDoFnTests
-        excludeCategories 'org.apache.beam.sdk.testing.UsesBoundedSplittableParDo'
+        // Unbounded SDF requires the streaming residual relay (#19468).
         excludeCategories 'org.apache.beam.sdk.testing.UsesUnboundedSplittableParDo'
         excludeCategories 'org.apache.beam.sdk.testing.UsesStrictTimerOrdering'
-        excludeCategories 'org.apache.beam.sdk.testing.UsesBundleFinalizer'
       }
       testFilter = {
         // TODO (https://github.com/apache/beam/issues/20189)
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
@@ -39,6 +39,7 @@
 import org.apache.beam.runners.spark.translation.SparkStreamingPortablePipelineTranslator;
 import org.apache.beam.runners.spark.translation.SparkStreamingTranslationContext;
 import org.apache.beam.runners.spark.translation.SparkTranslationContext;
+import org.apache.beam.runners.spark.translation.streaming.SdfResidualRelay;
 import org.apache.beam.runners.spark.util.GlobalWatermarkHolder;
 import org.apache.beam.sdk.io.FileSystems;
 import org.apache.beam.sdk.metrics.MetricsEnvironment;
@@ -166,8 +167,10 @@ public PortablePipelineResult run(RunnerApi.Pipeline pipeline, JobInfo jobInfo)
                   jssc.awaitTerminationOrTimeout(timeout);
                 } catch (InterruptedException e) {
                   LOG.warn("Streaming context interrupted, shutting down.", e);
+                } finally {
+                  jssc.stop();
+                  SdfResidualRelay.unregisterJob(jobInfo.jobId());
                 }
-                jssc.stop();
                 LOG.info("Job {} finished.", jobInfo.jobId());
               });
       result = new SparkPipelineResult.PortableStreamingMode(submissionFuture, jssc);
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkBatchPortablePipelineTranslator.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkBatchPortablePipelineTranslator.java
@@ -262,7 +262,9 @@ private static <InputT, OutputT, SideInputT> void translateExecutableStage(
               SparkExecutableStageContextFactory.getInstance(),
               broadcastVariables,
               MetricsAccumulator.getInstance(),
-              windowCoder);
+              windowCoder,
+              getWindowedValueCoder(inputPCollectionId, components),
+              false);
       staged = groupedByKey.flatMap(function.forPair());
     } else {
       JavaRDD<WindowedValue<InputT>> inputRdd2 = ((BoundedDataset<InputT>) inputDataset).getRDD();
@@ -275,7 +277,9 @@ private static <InputT, OutputT, SideInputT> void translateExecutableStage(
               SparkExecutableStageContextFactory.getInstance(),
               broadcastVariables,
               MetricsAccumulator.getInstance(),
-              windowCoder);
+              windowCoder,
+              getWindowedValueCoder(inputPCollectionId, components),
+              false);
       staged = inputRdd2.mapPartitions(function2);
     }
 
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkExecutableStageFunction.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkExecutableStageFunction.java
@@ -19,6 +19,7 @@
 
 import java.io.IOException;
 import java.io.Serializable;
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.EnumMap;
 import java.util.Iterator;
@@ -27,6 +28,7 @@
 import java.util.Map;
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.stream.Collectors;
+import org.apache.beam.model.fnexecution.v1.BeamFnApi.DelayedBundleApplication;
 import org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleProgressResponse;
 import org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleResponse;
 import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateKey;
@@ -36,6 +38,9 @@
 import org.apache.beam.runners.core.TimerInternals;
 import org.apache.beam.runners.core.construction.SerializablePipelineOptions;
 import org.apache.beam.runners.core.metrics.MetricsContainerImpl;
+import org.apache.beam.runners.fnexecution.control.BundleCheckpointHandler;
+import org.apache.beam.runners.fnexecution.control.BundleFinalizationHandlers;
+import org.apache.beam.runners.fnexecution.control.BundleFinalizationHandlers.InMemoryFinalizer;
 import org.apache.beam.runners.fnexecution.control.BundleProgressHandler;
 import org.apache.beam.runners.fnexecution.control.ExecutableStageContext;
 import org.apache.beam.runners.fnexecution.control.JobBundleFactory;
@@ -59,6 +64,7 @@
 import org.apache.beam.sdk.io.FileSystems;
 import org.apache.beam.sdk.transforms.join.RawUnionValue;
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.util.CoderUtils;
 import org.apache.beam.sdk.util.construction.Timer;
 import org.apache.beam.sdk.util.construction.graph.ExecutableStage;
 import org.apache.beam.sdk.values.WindowedValue;
@@ -85,6 +91,10 @@
 class SparkExecutableStageFunction<InputT, SideInputT>
     implements FlatMapFunction<Iterator<WindowedValue<InputT>>, RawUnionValue> {
 
+  // Union tag carrying serialized SDF residuals to the streaming residual relay. Regular output
+  // tags start at 0.
+  static final int SDF_RESIDUAL_TAG = -1;
+
   // Pipeline options for initializing the FileSystems
   private final SerializablePipelineOptions pipelineOptions;
   private final RunnerApi.ExecutableStagePayload stagePayload;
@@ -95,6 +105,10 @@ class SparkExecutableStageFunction<InputT, SideInputT>
       sideInputs;
   private final MetricsContainerStepMapAccumulator metricsAccumulator;
   private final Coder windowCoder;
+  // Coder for the stage input, used to re-feed SDF self-checkpoint residuals.
+  private final Coder<WindowedValue<InputT>> inputCoder;
+  // Streaming emits residuals to the cross-batch relay; batch drains them in place.
+  private final boolean emitSdfResiduals;
   private final JobInfo jobInfo;
 
   private transient InMemoryBagUserStateFactory bagUserStateHandlerFactory;
@@ -108,7 +122,9 @@ class SparkExecutableStageFunction<InputT, SideInputT>
       SparkExecutableStageContextFactory contextFactory,
       Map<String, Tuple2<Broadcast<List<byte[]>>, WindowedValueCoder<SideInputT>>> sideInputs,
       MetricsContainerStepMapAccumulator metricsAccumulator,
-      Coder windowCoder) {
+      Coder windowCoder,
+      Coder<WindowedValue<InputT>> inputCoder,
+      boolean emitSdfResiduals) {
     this.pipelineOptions = pipelineOptions;
     this.stagePayload = stagePayload;
     this.jobInfo = jobInfo;
@@ -117,6 +133,8 @@ class SparkExecutableStageFunction<InputT, SideInputT>
     this.sideInputs = sideInputs;
     this.metricsAccumulator = metricsAccumulator;
     this.windowCoder = windowCoder;
+    this.inputCoder = inputCoder;
+    this.emitSdfResiduals = emitSdfResiduals;
   }
 
   /** Call the executable stage function on the values of a PairRDD, ignoring the key. */
@@ -146,7 +164,30 @@ public Iterator<RawUnionValue> call(Iterator<WindowedValue<InputT>> inputs) thro
                 executableStage, stageBundleFactory.getProcessBundleDescriptor());
         if (executableStage.getTimers().size() == 0) {
           ReceiverFactory receiverFactory = new ReceiverFactory(collector, outputMap);
-          processElements(stateRequestHandler, receiverFactory, null, stageBundleFactory, inputs);
+          ResidualCollector residualCollector = new ResidualCollector();
+          InMemoryFinalizer finalizer =
+              BundleFinalizationHandlers.inMemoryFinalizer(
+                  stageBundleFactory.getInstructionRequestHandler());
+          processElements(
+              stateRequestHandler,
+              receiverFactory,
+              null,
+              stageBundleFactory,
+              inputs,
+              residualCollector,
+              finalizer);
+          if (emitSdfResiduals) {
+            for (DelayedBundleApplication residual : residualCollector.drain()) {
+              collector.add(new RawUnionValue(SDF_RESIDUAL_TAG, residual.toByteArray()));
+            }
+          } else {
+            processResiduals(
+                stateRequestHandler,
+                receiverFactory,
+                stageBundleFactory,
+                residualCollector,
+                finalizer);
+          }
           return collector.iterator();
         }
         // Used with Batch, we know that all the data is available for this key. We can't use the
@@ -172,8 +213,18 @@ public Iterator<RawUnionValue> call(Iterator<WindowedValue<InputT>> inputs) thro
                 windowCoder);
 
         // Process inputs.
+        ResidualCollector residualCollector = new ResidualCollector();
+        InMemoryFinalizer finalizer =
+            BundleFinalizationHandlers.inMemoryFinalizer(
+                stageBundleFactory.getInstructionRequestHandler());
         processElements(
-            stateRequestHandler, receiverFactory, timerReceiverFactory, stageBundleFactory, inputs);
+            stateRequestHandler,
+            receiverFactory,
+            timerReceiverFactory,
+            stageBundleFactory,
+            inputs,
+            residualCollector,
+            finalizer);
 
         // Finish any pending windows by advancing the input watermark to infinity.
         timerInternals.advanceInputWatermark(BoundedWindow.TIMESTAMP_MAX_VALUE);
@@ -189,12 +240,17 @@ public Iterator<RawUnionValue> call(Iterator<WindowedValue<InputT>> inputs) thro
                   receiverFactory,
                   timerReceiverFactory,
                   stateRequestHandler,
-                  getBundleProgressHandler())) {
+                  getBundleProgressHandler(),
+                  finalizer,
+                  residualCollector)) {
 
             PipelineTranslatorUtils.fireEligibleTimers(
                 timerInternals, bundle.getTimerReceivers(), currentTimerKey);
           }
+          finalizer.finalizeAllOutstandingBundles();
         }
+        processResiduals(
+            stateRequestHandler, receiverFactory, stageBundleFactory, residualCollector, finalizer);
         return collector.iterator();
       }
     }
@@ -207,21 +263,116 @@ private void processElements(
       ReceiverFactory receiverFactory,
       TimerReceiverFactory timerReceiverFactory,
       StageBundleFactory stageBundleFactory,
-      Iterator<WindowedValue<InputT>> inputs)
+      Iterator<WindowedValue<InputT>> inputs,
+      BundleCheckpointHandler checkpointHandler,
+      InMemoryFinalizer finalizer)
       throws Exception {
     try (RemoteBundle bundle =
         stageBundleFactory.getBundle(
             receiverFactory,
             timerReceiverFactory,
             stateRequestHandler,
-            getBundleProgressHandler())) {
+            getBundleProgressHandler(),
+            finalizer,
+            checkpointHandler)) {
       FnDataReceiver<WindowedValue<?>> mainReceiver =
           Iterables.getOnlyElement(bundle.getInputReceivers().values());
       while (inputs.hasNext()) {
         WindowedValue<InputT> input = inputs.next();
         mainReceiver.accept(input);
       }
     }
+    finalizer.finalizeAllOutstandingBundles();
+  }
+
+  // Re-feeds SDF self-checkpoint residuals in fresh bundles until the SDK returns none. Each
+  // residual resumes at its own requested time. Bounded restrictions always run out; unbounded ones
+  // never would, so they are rejected rather than drained forever.
+  private void processResiduals(
+      StateRequestHandler stateRequestHandler,
+      ReceiverFactory receiverFactory,
+      StageBundleFactory stageBundleFactory,
+      ResidualCollector residualCollector,
+      InMemoryFinalizer finalizer)
+      throws Exception {
+    List<ScheduledResidual> scheduled = schedule(residualCollector.drain());
+    while (!scheduled.isEmpty()) {
+      List<ScheduledResidual> due = takeDue(scheduled);
+      try (RemoteBundle bundle =
+          stageBundleFactory.getBundle(
+              receiverFactory,
+              null,
+              stateRequestHandler,
+              getBundleProgressHandler(),
+              finalizer,
+              residualCollector)) {
+        FnDataReceiver<WindowedValue<?>> mainReceiver =
+            Iterables.getOnlyElement(bundle.getInputReceivers().values());
+        for (ScheduledResidual residual : due) {
+          mainReceiver.accept(CoderUtils.decodeFromByteArray(inputCoder, residual.elementBytes));
+        }
+      }
+      finalizer.finalizeAllOutstandingBundles();
+      scheduled.addAll(schedule(residualCollector.drain()));
+    }
+  }
+
+  private static List<ScheduledResidual> schedule(List<DelayedBundleApplication> residuals) {
+    long now = System.currentTimeMillis();
+    List<ScheduledResidual> scheduled = new ArrayList<>();
+    for (DelayedBundleApplication residual : residuals) {
+      if (residual.getApplication().getElement().isEmpty()) {
+        continue;
+      }
+      if (residual.getApplication().getIsBounded() == RunnerApi.IsBounded.Enum.UNBOUNDED) {
+        throw new UnsupportedOperationException(
+            "Unbounded splittable DoFn is not supported in batch mode on the Spark runner. See "
+                + "https://github.com/apache/beam/issues/19468.");
+      }
+      long delayMillis =
+          residual.hasRequestedTimeDelay()
+              ? residual.getRequestedTimeDelay().getSeconds() * 1000
+                  + residual.getRequestedTimeDelay().getNanos() / 1_000_000
+              : 0;
+      scheduled.add(
+          new ScheduledResidual(
+              now + delayMillis, residual.getApplication().getElement().toByteArray()));
+    }
+    return scheduled;
+  }
+
+  // Waits for the earliest resume time, then removes and returns everything due by then.
+  private static List<ScheduledResidual> takeDue(List<ScheduledResidual> scheduled)
+      throws InterruptedException {
+    long earliest = Long.MAX_VALUE;
+    for (ScheduledResidual residual : scheduled) {
+      earliest = Math.min(earliest, residual.dueMillis);
+    }
+    long waitMillis = earliest - System.currentTimeMillis();
+    if (waitMillis > 0) {
+      Thread.sleep(waitMillis);
+    }
+    long now = System.currentTimeMillis();
+    List<ScheduledResidual> due = new ArrayList<>();
+    Iterator<ScheduledResidual> iterator = scheduled.iterator();
+    while (iterator.hasNext()) {
+      ScheduledResidual residual = iterator.next();
+      if (residual.dueMillis <= now) {
+        due.add(residual);
+        iterator.remove();
+      }
+    }
+    return due;
+  }
+
+  private static class ScheduledResidual {
+    private final long dueMillis;
+    private final byte[] elementBytes;
+
+    ScheduledResidual(long dueMillis, byte[] elementBytes) {
+      this.dueMillis = dueMillis;
+      this.elementBytes = elementBytes;
+    }
   }
 
   private BundleProgressHandler getBundleProgressHandler() {
@@ -295,6 +446,27 @@ interface JobBundleFactoryCreator extends Serializable {
     JobBundleFactory create();
   }
 
+  /** Collects SDF self-checkpoint residuals returned by the SDK harness. */
+  private static class ResidualCollector implements BundleCheckpointHandler {
+
+    private final ConcurrentLinkedQueue<DelayedBundleApplication> residuals =
+        new ConcurrentLinkedQueue<>();
+
+    @Override
+    public void onCheckpoint(ProcessBundleResponse response) {
+      residuals.addAll(response.getResidualRootsList());
+    }
+
+    private List<DelayedBundleApplication> drain() {
+      List<DelayedBundleApplication> pending = new ArrayList<>();
+      DelayedBundleApplication residual;
+      while ((residual = residuals.poll()) != null) {
+        pending.add(residual);
+      }
+      return pending;
+    }
+  }
+
   /**
    * Receiver factory that wraps outgoing elements with the corresponding union tag for a
    * multiplexed PCollection.
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
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/streaming/SdfResidualRelay.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/streaming/SdfResidualRelay.java
diff --git a/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/SparkExecutableStageFunctionTest.java b/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/SparkExecutableStageFunctionTest.java
diff --git a/sdks/python/apache_beam/runners/portability/spark_runner_test.py b/sdks/python/apache_beam/runners/portability/spark_runner_test.py

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`	`1`	`{`
`2`	`2`	`"comment": "Modify this file in a trivial way to cause this test suite to run",`
`3`		`- "modification": 6`
	`3`	`+ "modification": 7`
`4`	`4`	`}`
Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`	`1`	`{`
`2`	`2`	`"comment": "Modify this file in a trivial way to cause this test suite to run",`
`3`		`- "modification": 1`
	`3`	`+ "modification": 2`
`4`	`4`	`}`
Original file line number	Diff line number	Diff line change
`@@ -2,5 +2,6 @@`
`2`	`2`	`"https://github.com/apache/beam/pull/34830": "testing",`
`3`	`3`	`"https://github.com/apache/beam/issues/35429": "testing",`
`4`	`4`	`"trigger-2026-04-04": "portable_runner expand_sdf opt-in",`
`5`		`- "https://github.com/apache/beam/pull/38892": "UnboundedSource portable VR test"`
	`5`	`+ "https://github.com/apache/beam/pull/38892": "UnboundedSource portable VR test",`
	`6`	`+ "https://github.com/apache/beam/issues/19468": "SDF self-checkpointing and bundle finalization"`
`6`	`7`	`}`