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

.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
 }

.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
 }

.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
 }

.github/trigger_files/beam_PostCommit_Python_ValidatesRunner_Spark.json

@@ -3,5 +3,6 @@
   "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",
-  "modification": 1
+  "modification": 1,
+  "https://github.com/apache/beam/issues/19468": "SDF self-checkpointing and bundle finalization"
 }

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

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)

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);

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);
     }
 

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.