This is an automated email from the ASF dual-hosted git repository.
Abacn pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new f9ca09b3183 [Spark] Support splittable DoFn self-checkpointing in
portable batch (#39331)
f9ca09b3183 is described below
commit f9ca09b31834a367d671cb76397734c7b1ea3b92
Author: Elia Liu <[email protected]>
AuthorDate: Sat Aug 8 01:03:17 2026 +1000
[Spark] Support splittable DoFn self-checkpointing in portable batch
(#39331)
* [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.
---
.../beam_PostCommit_Java_PVR_Spark3_Streaming.json | 2 +-
.../beam_PostCommit_Java_PVR_Spark_Batch.json | 2 +-
...beam_PostCommit_Java_ValidatesRunner_Spark.json | 2 +-
...am_PostCommit_Python_ValidatesRunner_Spark.json | 3 +-
runners/spark/job-server/spark_job_server.gradle | 5 +-
.../SparkBatchPortablePipelineTranslator.java | 8 +-
.../translation/SparkExecutableStageFunction.java | 172 +++++++++++++++++++--
.../SparkStreamingPortablePipelineTranslator.java | 4 +-
.../SparkExecutableStageFunctionTest.java | 138 ++++++++++++++++-
.../runners/portability/spark_runner_test.py | 20 ---
10 files changed, 307 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 455144f02a3..d6a91b7e2e8 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 455144f02a3..d6a91b7e2e8 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 1efc8e9e440..3f63c0c9975 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 f4ec72dc416..6384446f50e 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/runners/spark/job-server/spark_job_server.gradle
b/runners/spark/job-server/spark_job_server.gradle
index 5240bb310d0..2811f875f84 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 ba3aa0e4d24..521f9583597 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 @@ public class SparkBatchPortablePipelineTranslator
SparkExecutableStageContextFactory.getInstance(),
broadcastVariables,
MetricsAccumulator.getInstance(),
- windowCoder);
+ windowCoder,
+ getWindowedValueCoder(inputPCollectionId, components),
+ true);
staged = groupedByKey.flatMap(function.forPair());
} else {
JavaRDD<WindowedValue<InputT>> inputRdd2 = ((BoundedDataset<InputT>)
inputDataset).getRDD();
@@ -275,7 +277,9 @@ public class SparkBatchPortablePipelineTranslator
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 757740e2df5..a04899bf9c9 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 @@ package org.apache.beam.runners.spark.translation;
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.ProcessBundleResponse;
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;
@@ -57,8 +63,10 @@ import org.apache.beam.runners.spark.util.ByteArray;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.fn.data.FnDataReceiver;
import org.apache.beam.sdk.io.FileSystems;
+import org.apache.beam.sdk.state.MapState;
import org.apache.beam.sdk.transforms.join.RawUnionValue;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+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<InputT, SideInputT>
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<WindowedValue<InputT>> 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<InputT, SideInputT>
SparkExecutableStageContextFactory contextFactory,
Map<String, Tuple2<Broadcast<List<byte[]>>,
WindowedValueCoder<SideInputT>>> sideInputs,
MetricsContainerStepMapAccumulator metricsAccumulator,
- Coder windowCoder) {
+ Coder windowCoder,
+ Coder<WindowedValue<InputT>> inputCoder,
+ boolean batch) {
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.batch = batch;
}
/** Call the executable stage function on the values of a PairRDD, ignoring
the key. */
@@ -144,9 +162,18 @@ class SparkExecutableStageFunction<InputT, SideInputT>
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 +200,12 @@ class SparkExecutableStageFunction<InputT, SideInputT>
// 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 +214,30 @@ class SparkExecutableStageFunction<InputT, SideInputT>
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 +250,17 @@ class SparkExecutableStageFunction<InputT, SideInputT>
ReceiverFactory receiverFactory,
TimerReceiverFactory timerReceiverFactory,
StageBundleFactory stageBundleFactory,
- Iterator<WindowedValue<InputT>> inputs)
+ Iterator<WindowedValue<InputT>> inputs,
+ BundleCheckpointHandler checkpointHandler)
throws Exception {
try (RemoteBundle bundle =
stageBundleFactory.getBundle(
receiverFactory,
timerReceiverFactory,
stateRequestHandler,
- getBundleProgressHandler())) {
+ getBundleProgressHandler(),
+ getBundleFinalizationHandler(),
+ checkpointHandler)) {
FnDataReceiver<WindowedValue<?>> mainReceiver =
Iterables.getOnlyElement(bundle.getInputReceivers().values());
while (inputs.hasNext()) {
@@ -224,6 +270,100 @@ class SparkExecutableStageFunction<InputT, SideInputT>
}
}
+ 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<WindowedValue<InputT>> residuals = new ArrayList<>();
+ TimerInternals.TimerData timer;
+ while ((timer = sdfTimerInternals.removeNextProcessingTimer()) !=
null) {
+ MapState<String, WindowedValue<InputT>> residualState =
+ sdfStateInternals.state(
+ timer.getNamespace(),
+
BundleCheckpointHandlers.StateAndTimerBundleCheckpointHandler.residualStateTag(
+ inputCoder));
+ WindowedValue<InputT> residual =
residualState.get(timer.getTimerId()).read();
+ residualState.remove(timer.getTimerId());
+ if (residual != null) {
+ residuals.add(residual);
+ }
+ }
+ FnDataReceiver<WindowedValue<?>> mainReceiver =
+ Iterables.getOnlyElement(bundle.getInputReceivers().values());
+ for (WindowedValue<InputT> 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 9975c81b56a..db3551454dc 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 @@ public class SparkStreamingPortablePipelineTranslator
SparkExecutableStageContextFactory.getInstance(),
broadcastVariables,
MetricsAccumulator.getInstance(),
- windowCoder);
+ windowCoder,
+ getWindowedValueCoder(inputPCollectionId, components),
+ false);
JavaDStream<RawUnionValue> 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 98601389f5c..31ba649cfff 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;
@@ -33,6 +34,9 @@ import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
+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.model.pipeline.v1.RunnerApi;
import org.apache.beam.model.pipeline.v1.RunnerApi.Components;
import org.apache.beam.model.pipeline.v1.RunnerApi.ExecutableStagePayload;
@@ -52,13 +56,18 @@ import
org.apache.beam.runners.fnexecution.control.StageBundleFactory;
import org.apache.beam.runners.fnexecution.control.TimerReceiverFactory;
import org.apache.beam.runners.fnexecution.state.StateRequestHandler;
import
org.apache.beam.runners.spark.metrics.MetricsContainerStepMapAccumulator;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.VarIntCoder;
import org.apache.beam.sdk.fn.data.FnDataReceiver;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.transforms.join.RawUnionValue;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.util.CoderUtils;
import org.apache.beam.sdk.util.construction.Timer;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.WindowedValue;
import org.apache.beam.sdk.values.WindowedValues;
+import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
import org.junit.Before;
import org.junit.Test;
@@ -97,12 +106,31 @@ 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<String, FnDataReceiver> inputReceiver =
@@ -126,7 +154,8 @@ public class SparkExecutableStageFunctionTest {
SparkExecutableStageFunction<Integer, ?> 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 +276,8 @@ public class SparkExecutableStageFunctionTest {
List<WindowedValue<Integer>> 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 +290,104 @@ public class SparkExecutableStageFunctionTest {
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<WindowedValue<?>> received = new ArrayList<>();
+ WindowedValue<Integer> residualValue =
WindowedValues.valueInGlobalWindow(7);
+ Coder<WindowedValue<Integer>> 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<String, FnDataReceiver> getInputReceivers() {
+ FnDataReceiver<WindowedValue<?>> receiver = received::add;
+ return ImmutableMap.of("input", receiver);
+ }
+
+ @Override
+ public Map<KV<String, String>, FnDataReceiver<Timer>>
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<Integer, ?> 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 <InputT, SideInputT> SparkExecutableStageFunction<InputT,
SideInputT> getFunction(
Map<String, Integer> outputMap) {
return new SparkExecutableStageFunction<>(
@@ -270,6 +398,8 @@ public class SparkExecutableStageFunctionTest {
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 4152b8d09f4..40774eb9602 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 @@ class
SparkRunnerTest(portable_runner_test.PortableRunnerTest):
# 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 @@ class
SparkRunnerTest(portable_runner_test.PortableRunnerTest):
# 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";)
