From 6d6110542d68f457be3227daeef38542a30acd00 Mon Sep 17 00:00:00 2001 From: Eliaazzz Date: Wed, 15 Jul 2026 12:04:01 +1000 Subject: [PATCH] [Spark] Support splittable DoFn self-checkpointing in portable batch The portable Spark runner never passed a BundleCheckpointHandler to StageBundleFactory.getBundle, so a splittable DoFn that self-checkpoints failed on its first bundle and could not run at all. In batch, a stage containing a splittable DoFn now holds each residual in memory under a processing time timer, the way the portable Flink batch runner does. Once the stage has drained its inputs, processing time advances to infinity and the held residuals are replayed until the SDK stops asking to resume, so a bounded restriction always runs out. Streaming keeps rejecting self-checkpointing, with a message naming the issue, since a residual has nowhere to live across micro-batches. Bundle finalization is likewise rejected rather than run early, since this runner cannot report that a bundle's output is durably committed. Unskips the bounded splittable DoFn tests for the Spark runner. --- ..._PostCommit_Java_PVR_Spark3_Streaming.json | 2 +- .../beam_PostCommit_Java_PVR_Spark_Batch.json | 2 +- ...PostCommit_Java_ValidatesRunner_Spark.json | 2 +- ...stCommit_Python_ValidatesRunner_Spark.json | 3 +- CHANGES.md | 1 + .../spark/job-server/spark_job_server.gradle | 5 +- .../SparkBatchPortablePipelineTranslator.java | 8 +- .../SparkExecutableStageFunction.java | 167 ++++++++++++++++-- ...rkStreamingPortablePipelineTranslator.java | 4 +- .../SparkExecutableStageFunctionTest.java | 152 +++++++++++++++- .../runners/portability/spark_runner_test.py | 20 --- 11 files changed, 317 insertions(+), 49 deletions(-) diff --git a/.github/trigger_files/beam_PostCommit_Java_PVR_Spark3_Streaming.json b/.github/trigger_files/beam_PostCommit_Java_PVR_Spark3_Streaming.json index 455144f02a35..d6a91b7e2e86 100644 --- 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 index 455144f02a35..d6a91b7e2e86 100644 --- 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 index 1efc8e9e4405..3f63c0c9975f 100644 --- 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 index f4ec72dc416b..6384446f50e4 100644 --- a/.github/trigger_files/beam_PostCommit_Python_ValidatesRunner_Spark.json +++ b/.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" } diff --git a/CHANGES.md b/CHANGES.md index d853314a0ad3..50edf85f3494 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -76,6 +76,7 @@ * (Python) Added `equal_to_approx`, an `assert_that` matcher that compares numeric pipeline outputs with a configurable tolerance ([#18028](https://github.com/apache/beam/issues/18028)). * (Python) `Timestamp` now supports variable subsecond precision, up to nanoseconds. The portable `beam:logical_type:timestamp:v1` logical type now maps to Python's `Timestamp` ([#39344](https://github.com/apache/beam/issues/39344)). +* Splittable DoFn self-checkpointing is now supported on the portable Spark runner in batch mode, for bounded restrictions ([#19468](https://github.com/apache/beam/issues/19468)). * 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 index 5240bb310d05..2811f875f84f 100644 --- a/runners/spark/job-server/spark_job_server.gradle +++ b/runners/spark/job-server/spark_job_server.gradle @@ -199,10 +199,11 @@ 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' + // TODO (https://github.com/apache/beam/issues/19468) unbounded SDF needs residuals to + // survive across micro-batches, which the streaming path cannot do yet. excludeCategories 'org.apache.beam.sdk.testing.UsesUnboundedSplittableParDo' excludeCategories 'org.apache.beam.sdk.testing.UsesStrictTimerOrdering' + // TODO (https://github.com/apache/beam/issues/19517) bundle finalization excludeCategories 'org.apache.beam.sdk.testing.UsesBundleFinalizer' } testFilter = { 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 index ba3aa0e4d24a..521f95835978 100644 --- 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 void translateExecutableStage( SparkExecutableStageContextFactory.getInstance(), broadcastVariables, MetricsAccumulator.getInstance(), - windowCoder); + windowCoder, + getWindowedValueCoder(inputPCollectionId, components), + true); staged = groupedByKey.flatMap(function.forPair()); } else { JavaRDD> inputRdd2 = ((BoundedDataset) inputDataset).getRDD(); @@ -275,7 +277,9 @@ private static void translateExecutableStage( SparkExecutableStageContextFactory.getInstance(), broadcastVariables, MetricsAccumulator.getInstance(), - windowCoder); + windowCoder, + getWindowedValueCoder(inputPCollectionId, components), + true); 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 index 757740e2df5a..7744548df0aa 100644 --- 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; @@ -32,10 +33,15 @@ import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateKey; import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateKey.TypeCase; import org.apache.beam.model.pipeline.v1.RunnerApi; +import org.apache.beam.runners.core.InMemoryStateInternals; import org.apache.beam.runners.core.InMemoryTimerInternals; +import org.apache.beam.runners.core.StateInternals; 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.BundleCheckpointHandlers; +import org.apache.beam.runners.fnexecution.control.BundleFinalizationHandler; 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 +65,8 @@ 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.state.MapState; +import org.apache.beam.sdk.util.construction.PTransformTranslation; import org.apache.beam.sdk.util.construction.Timer; import org.apache.beam.sdk.util.construction.graph.ExecutableStage; import org.apache.beam.sdk.values.WindowedValue; @@ -95,10 +103,16 @@ class SparkExecutableStageFunction sideInputs; private final MetricsContainerStepMapAccumulator metricsAccumulator; private final Coder windowCoder; + // Coder for this stage's input, used to hold and replay splittable DoFn residuals. + private final Coder> inputCoder; + // Batch replays residuals in place. Streaming has nowhere to hold them across micro-batches yet. + private final boolean batch; private final JobInfo jobInfo; private transient InMemoryBagUserStateFactory bagUserStateHandlerFactory; private transient Object currentTimerKey; + private transient InMemoryTimerInternals sdfTimerInternals; + private transient StateInternals sdfStateInternals; SparkExecutableStageFunction( SerializablePipelineOptions pipelineOptions, @@ -108,7 +122,9 @@ class SparkExecutableStageFunction SparkExecutableStageContextFactory contextFactory, Map>, WindowedValueCoder>> sideInputs, MetricsContainerStepMapAccumulator metricsAccumulator, - Coder windowCoder) { + Coder windowCoder, + Coder> inputCoder, + boolean batch) { this.pipelineOptions = pipelineOptions; this.stagePayload = stagePayload; this.jobInfo = jobInfo; @@ -117,6 +133,8 @@ class SparkExecutableStageFunction this.sideInputs = sideInputs; this.metricsAccumulator = metricsAccumulator; this.windowCoder = windowCoder; + this.inputCoder = inputCoder; + this.batch = batch; } /** Call the executable stage function on the values of a PairRDD, ignoring the key. */ @@ -144,9 +162,13 @@ public Iterator call(Iterator> inputs) thro StateRequestHandler stateRequestHandler = getStateRequestHandler( executableStage, stageBundleFactory.getProcessBundleDescriptor()); + BundleCheckpointHandler checkpointHandler = getBundleCheckpointHandler(executableStage); if (executableStage.getTimers().size() == 0) { ReceiverFactory receiverFactory = new ReceiverFactory(collector, outputMap); - processElements(stateRequestHandler, receiverFactory, null, stageBundleFactory, inputs); + processElements( + stateRequestHandler, receiverFactory, null, stageBundleFactory, inputs, checkpointHandler); + replaySdfResiduals( + stateRequestHandler, receiverFactory, null, stageBundleFactory, checkpointHandler); return collector.iterator(); } // Used with Batch, we know that all the data is available for this key. We can't use the @@ -173,7 +195,12 @@ public Iterator call(Iterator> inputs) thro // Process inputs. processElements( - stateRequestHandler, receiverFactory, timerReceiverFactory, stageBundleFactory, inputs); + stateRequestHandler, + receiverFactory, + timerReceiverFactory, + stageBundleFactory, + inputs, + checkpointHandler); // Finish any pending windows by advancing the input watermark to infinity. timerInternals.advanceInputWatermark(BoundedWindow.TIMESTAMP_MAX_VALUE); @@ -182,19 +209,30 @@ public Iterator call(Iterator> inputs) thro timerInternals.advanceSynchronizedProcessingTime(BoundedWindow.TIMESTAMP_MAX_VALUE); // Now we fire the timers and process elements generated by timers (which may be timers - // itself) - while (timerInternals.hasPendingTimers()) { - try (RemoteBundle bundle = - stageBundleFactory.getBundle( - receiverFactory, - timerReceiverFactory, - stateRequestHandler, - getBundleProgressHandler())) { + // itself). A replayed splittable DoFn residual can set a timer, and a fired timer can + // produce a residual, so alternate until neither has anything left. + do { + while (timerInternals.hasPendingTimers()) { + try (RemoteBundle bundle = + stageBundleFactory.getBundle( + receiverFactory, + timerReceiverFactory, + stateRequestHandler, + getBundleProgressHandler(), + getBundleFinalizationHandler(), + checkpointHandler)) { - PipelineTranslatorUtils.fireEligibleTimers( - timerInternals, bundle.getTimerReceivers(), currentTimerKey); + PipelineTranslatorUtils.fireEligibleTimers( + timerInternals, bundle.getTimerReceivers(), currentTimerKey); + } } - } + replaySdfResiduals( + stateRequestHandler, + receiverFactory, + timerReceiverFactory, + stageBundleFactory, + checkpointHandler); + } while (timerInternals.hasPendingTimers()); return collector.iterator(); } } @@ -207,14 +245,17 @@ private void processElements( ReceiverFactory receiverFactory, TimerReceiverFactory timerReceiverFactory, StageBundleFactory stageBundleFactory, - Iterator> inputs) + Iterator> inputs, + BundleCheckpointHandler checkpointHandler) throws Exception { try (RemoteBundle bundle = stageBundleFactory.getBundle( receiverFactory, timerReceiverFactory, stateRequestHandler, - getBundleProgressHandler())) { + getBundleProgressHandler(), + getBundleFinalizationHandler(), + checkpointHandler)) { FnDataReceiver> mainReceiver = Iterables.getOnlyElement(bundle.getInputReceivers().values()); while (inputs.hasNext()) { @@ -224,6 +265,100 @@ private void processElements( } } + private static boolean hasSdf(ExecutableStage executableStage) { + return executableStage.getTransforms().stream() + .anyMatch( + transform -> + transform + .getTransform() + .getSpec() + .getUrn() + .equals( + PTransformTranslation + .SPLITTABLE_PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS_URN)); + } + + // Holds a splittable DoFn's self-checkpoint residual in memory under a processing time timer, so + // it can be replayed once this stage has drained its inputs. + private BundleCheckpointHandler getBundleCheckpointHandler(ExecutableStage executableStage) { + sdfTimerInternals = null; + sdfStateInternals = null; + if (!batch) { + return response -> { + throw new UnsupportedOperationException( + "Splittable DoFn self-checkpointing is not supported on the portable Spark runner in " + + "streaming mode. For more details, please refer to " + + "https://github.com/apache/beam/issues/19468."); + }; + } + if (!hasSdf(executableStage)) { + return response -> { + throw new UnsupportedOperationException( + "Self-checkpoint is only supported on splittable DoFn."); + }; + } + sdfTimerInternals = new InMemoryTimerInternals(); + sdfStateInternals = InMemoryStateInternals.forKey("sdf_state"); + return new BundleCheckpointHandlers.StateAndTimerBundleCheckpointHandler<>( + key -> sdfTimerInternals, key -> sdfStateInternals, inputCoder, windowCoder); + } + + // Bundle finalization needs the runner to have durably committed the bundle's output first, which + // this runner cannot report, so it is rejected rather than silently finalized early. + private BundleFinalizationHandler getBundleFinalizationHandler() { + return bundleId -> { + throw new UnsupportedOperationException( + "The portable Spark runner does not support bundle finalization. For more details, please " + + "refer to https://github.com/apache/beam/issues/19517."); + }; + } + + // Replays held residuals until the splittable DoFn stops asking to resume. Processing time is at + // infinity, so every residual is due immediately and a bounded restriction always runs out. + private void replaySdfResiduals( + StateRequestHandler stateRequestHandler, + ReceiverFactory receiverFactory, + TimerReceiverFactory timerReceiverFactory, + StageBundleFactory stageBundleFactory, + BundleCheckpointHandler checkpointHandler) + throws Exception { + if (sdfTimerInternals == null) { + return; + } + sdfTimerInternals.advanceProcessingTime(BoundedWindow.TIMESTAMP_MAX_VALUE); + sdfTimerInternals.advanceSynchronizedProcessingTime(BoundedWindow.TIMESTAMP_MAX_VALUE); + while (sdfTimerInternals.hasPendingTimers()) { + try (RemoteBundle bundle = + stageBundleFactory.getBundle( + receiverFactory, + timerReceiverFactory, + stateRequestHandler, + getBundleProgressHandler(), + getBundleFinalizationHandler(), + checkpointHandler)) { + List> residuals = new ArrayList<>(); + TimerInternals.TimerData timer; + while ((timer = sdfTimerInternals.removeNextProcessingTimer()) != null) { + MapState> residualState = + sdfStateInternals.state( + timer.getNamespace(), + BundleCheckpointHandlers.StateAndTimerBundleCheckpointHandler.residualStateTag( + inputCoder)); + WindowedValue residual = residualState.get(timer.getTimerId()).read(); + residualState.remove(timer.getTimerId()); + if (residual != null) { + residuals.add(residual); + } + } + FnDataReceiver> mainReceiver = + Iterables.getOnlyElement(bundle.getInputReceivers().values()); + for (WindowedValue residual : residuals) { + mainReceiver.accept(residual); + } + } + } + } + private BundleProgressHandler getBundleProgressHandler() { String stageName = stagePayload.getInput(); MetricsContainerImpl container = metricsAccumulator.value().getContainer(stageName); 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 9975c81b56a4..db3551454dcd 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 @@ -251,7 +251,9 @@ private static void translateExecutableStage( SparkExecutableStageContextFactory.getInstance(), broadcastVariables, MetricsAccumulator.getInstance(), - windowCoder); + windowCoder, + getWindowedValueCoder(inputPCollectionId, components), + false); JavaDStream staged = inputDStream.mapPartitions(function); String intermediateId = getExecutableStageIntermediateId(transformNode); 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 index 98601389f5c9..97ef29855428 100644 --- 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 @@ -18,6 +18,7 @@ package org.apache.beam.runners.spark.translation; import static org.apache.beam.sdk.util.construction.PTransformTranslation.PAR_DO_TRANSFORM_URN; +import static org.apache.beam.sdk.util.construction.PTransformTranslation.SPLITTABLE_PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS_URN; import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.contains; import static org.mockito.ArgumentMatchers.any; @@ -56,6 +57,14 @@ import org.apache.beam.sdk.options.PipelineOptionsFactory; import org.apache.beam.sdk.transforms.join.RawUnionValue; import org.apache.beam.sdk.util.construction.Timer; +import org.apache.beam.model.fnexecution.v1.BeamFnApi.BundleApplication; +import org.apache.beam.model.fnexecution.v1.BeamFnApi.DelayedBundleApplication; +import org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleResponse; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.VarIntCoder; +import org.apache.beam.sdk.transforms.windowing.GlobalWindow; +import org.apache.beam.sdk.util.CoderUtils; +import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.WindowedValue; import org.apache.beam.sdk.values.WindowedValues; @@ -97,12 +106,37 @@ public class SparkExecutableStageFunctionTest { .build()) .build(); + private final ExecutableStagePayload sdfStagePayload = + ExecutableStagePayload.newBuilder() + .setInput(inputId) + .addTransforms("sdf-transform-id") + .setComponents( + Components.newBuilder() + .putTransforms( + "sdf-transform-id", + RunnerApi.PTransform.newBuilder() + .putInputs("input-name", inputId) + .setSpec( + RunnerApi.FunctionSpec.newBuilder() + .setUrn( + SPLITTABLE_PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS_URN)) + .build()) + .putPcollections(inputId, PCollection.getDefaultInstance()) + .build()) + .build(); + @Before public void setUpMocks() throws Exception { MockitoAnnotations.initMocks(this); when(contextFactory.get(any())).thenReturn(stageContext); when(stageContext.getStageBundleFactory(any())).thenReturn(stageBundleFactory); - when(stageBundleFactory.getBundle(any(), any(), any(), any(BundleProgressHandler.class))) + when(stageBundleFactory.getBundle( + any(), + any(), + any(), + any(BundleProgressHandler.class), + any(), + any())) .thenReturn(remoteBundle); @SuppressWarnings("unchecked") ImmutableMap inputReceiver = @@ -126,7 +160,13 @@ public void expectedInputsAreSent() throws Exception { SparkExecutableStageFunction function = getFunction(Collections.emptyMap()); RemoteBundle bundle = Mockito.mock(RemoteBundle.class); - when(stageBundleFactory.getBundle(any(), any(), any(), any(BundleProgressHandler.class))) + when(stageBundleFactory.getBundle( + any(), + any(), + any(), + any(BundleProgressHandler.class), + any(), + any())) .thenReturn(bundle); @SuppressWarnings("unchecked") @@ -247,7 +287,9 @@ public void testStageBundleClosed() throws Exception { List> inputs = new ArrayList<>(); inputs.add(WindowedValues.valueInGlobalWindow(0)); function.call(inputs.iterator()); - verify(stageBundleFactory).getBundle(any(), any(), any(), any(BundleProgressHandler.class)); + verify(stageBundleFactory) + .getBundle( + any(), any(), any(), any(BundleProgressHandler.class), any(), any()); verify(stageBundleFactory).getProcessBundleDescriptor(); verify(stageBundleFactory).close(); verifyNoMoreInteractions(stageBundleFactory); @@ -260,6 +302,106 @@ public void testNoCallOnEmptyInputIterator() throws Exception { verifyNoInteractions(stageBundleFactory); } + @Test + public void sdfResidualsAreReplayedUntilDrained() throws Exception { + // A stage whose bundle self-checkpoints once: the first bundle returns a residual, the replay + // bundle returns none. + List> received = new ArrayList<>(); + WindowedValue residualValue = WindowedValues.valueInGlobalWindow(7); + Coder> residualCoder = + WindowedValues.getFullCoder(VarIntCoder.of(), GlobalWindow.Coder.INSTANCE); + ProcessBundleResponse withResidual = + ProcessBundleResponse.newBuilder() + .addResidualRoots( + DelayedBundleApplication.newBuilder() + .setApplication( + BundleApplication.newBuilder() + .setElement( + ByteString.copyFrom( + CoderUtils.encodeToByteArray(residualCoder, residualValue))))) + .build(); + + StageBundleFactory bundleFactory = + new StageBundleFactory() { + private int bundles; + + @Override + public RemoteBundle getBundle( + OutputReceiverFactory receiverFactory, + TimerReceiverFactory timerReceiverFactory, + StateRequestHandler stateRequestHandler, + BundleProgressHandler progressHandler, + BundleFinalizationHandler finalizationHandler, + BundleCheckpointHandler checkpointHandler) { + boolean checkpointThisBundle = bundles++ == 0; + return new RemoteBundle() { + @Override + public String getId() { + return "bundle-id"; + } + + @Override + public Map getInputReceivers() { + FnDataReceiver> receiver = received::add; + return ImmutableMap.of("input", receiver); + } + + @Override + public Map, FnDataReceiver> getTimerReceivers() { + return Collections.emptyMap(); + } + + @Override + public void requestProgress() {} + + @Override + public void split(double fractionOfRemainder) {} + + @Override + public void close() { + if (checkpointThisBundle) { + checkpointHandler.onCheckpoint(withResidual); + } + } + }; + } + + @Override + public ProcessBundleDescriptors.ExecutableProcessBundleDescriptor + getProcessBundleDescriptor() { + return null; + } + + @Override + public InstructionRequestHandler getInstructionRequestHandler() { + return null; + } + + @Override + public void close() {} + }; + when(stageContext.getStageBundleFactory(any())).thenReturn(bundleFactory); + + SparkExecutableStageFunction function = + new SparkExecutableStageFunction<>( + pipelineOptions, + sdfStagePayload, + null, + Collections.emptyMap(), + contextFactory, + Collections.emptyMap(), + metricsAccumulator, + GlobalWindow.Coder.INSTANCE, + residualCoder, + true); + + function.call(Collections.singletonList(WindowedValues.valueInGlobalWindow(1)).iterator()); + + assertThat( + received, + contains(WindowedValues.valueInGlobalWindow(1), residualValue)); + } + private SparkExecutableStageFunction getFunction( Map outputMap) { return new SparkExecutableStageFunction<>( @@ -270,6 +412,8 @@ private SparkExecutableStageFunction ge contextFactory, Collections.emptyMap(), metricsAccumulator, - null); + null, + null, + true); } } diff --git a/sdks/python/apache_beam/runners/portability/spark_runner_test.py b/sdks/python/apache_beam/runners/portability/spark_runner_test.py index 4152b8d09f4f..40774eb9602c 100644 --- a/sdks/python/apache_beam/runners/portability/spark_runner_test.py +++ b/sdks/python/apache_beam/runners/portability/spark_runner_test.py @@ -144,26 +144,10 @@ def test_metrics(self): # Skip until Spark runner supports metrics. raise unittest.SkipTest("https://github.com/apache/beam/issues/19496") - def test_sdf(self): - # Skip until Spark runner supports SDF. - raise unittest.SkipTest("https://github.com/apache/beam/issues/19468") - def test_unbounded_source_read(self): # Skip until Spark runner supports SDF. raise unittest.SkipTest("https://github.com/apache/beam/issues/19468") - def test_sdf_with_watermark_tracking(self): - # Skip until Spark runner supports SDF. - raise unittest.SkipTest("https://github.com/apache/beam/issues/19468") - - def test_sdf_with_sdf_initiated_checkpointing(self): - # Skip until Spark runner supports SDF. - raise unittest.SkipTest("https://github.com/apache/beam/issues/19468") - - def test_sdf_synthetic_source(self): - # Skip until Spark runner supports SDF. - raise unittest.SkipTest("https://github.com/apache/beam/issues/19468") - def test_callbacks_with_exception(self): # Skip until Spark runner supports bundle finalization. raise unittest.SkipTest("https://github.com/apache/beam/issues/19517") @@ -172,10 +156,6 @@ def test_register_finalizations(self): # Skip until Spark runner supports bundle finalization. raise unittest.SkipTest("https://github.com/apache/beam/issues/19517") - def test_sdf_with_dofn_as_watermark_estimator(self): - # Skip until Spark runner supports SDF and self-checkpoint. - raise unittest.SkipTest("https://github.com/apache/beam/issues/19468") - def test_pardo_dynamic_timer(self): raise unittest.SkipTest("https://github.com/apache/beam/issues/20179")