You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by ie...@apache.org on 2019/04/03 14:16:24 UTC
[beam] branch master updated: [BEAM-2590] Implement basic Spark
portable runner
This is an automated email from the ASF dual-hosted git repository.
iemejia 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 5c7fc67 [BEAM-2590] Implement basic Spark portable runner
new 67425bf Merge pull request #8115: [BEAM-2590] Implement basic Spark portable runner
5c7fc67 is described below
commit 5c7fc67d4b2099aeb4951391bbc5ea54c8a0e2dc
Author: Kyle Weaver <kc...@google.com>
AuthorDate: Thu Mar 21 09:51:32 2019 -0700
[BEAM-2590] Implement basic Spark portable runner
---
runners/spark/build.gradle | 3 +
.../apache/beam/runners/spark/SparkJobInvoker.java | 78 ++++++++
.../beam/runners/spark/SparkJobServerDriver.java | 86 +++++++++
.../beam/runners/spark/SparkPipelineRunner.java | 79 ++++++++
.../runners/spark/translation/BoundedDataset.java | 2 +-
.../SparkBatchPortablePipelineTranslator.java | 210 +++++++++++++++++++++
.../SparkExecutableStageExtractionFunction.java | 42 +++++
.../translation/SparkExecutableStageFunction.java | 155 +++++++++++++++
.../spark/translation/SparkTranslationContext.java | 71 +++++++
.../runners/spark/SparkPortableExecutionTest.java | 139 ++++++++++++++
.../SparkExecutableStageFunctionTest.java | 203 ++++++++++++++++++++
11 files changed, 1067 insertions(+), 1 deletion(-)
diff --git a/runners/spark/build.gradle b/runners/spark/build.gradle
index 9dcdee8..f37a9d0 100644
--- a/runners/spark/build.gradle
+++ b/runners/spark/build.gradle
@@ -58,10 +58,12 @@ dependencies {
shadow project(path: ":beam-sdks-java-core", configuration: "shadow")
shadow project(path: ":beam-runners-core-construction-java", configuration: "shadow")
shadow project(path: ":beam-runners-core-java", configuration: "shadow")
+ shadow project(path: ":beam-runners-java-fn-execution", configuration: "shadow")
shadow library.java.guava
shadow library.java.jackson_annotations
shadow library.java.slf4j_api
shadow library.java.joda_time
+ shadow library.java.args4j
shadow "io.dropwizard.metrics:metrics-core:3.1.2"
shadow library.java.jackson_module_scala
provided library.java.spark_core
@@ -81,6 +83,7 @@ dependencies {
shadowTest project(path: ":beam-sdks-java-core", configuration: "shadowTest")
// SparkStateInternalsTest extends abstract StateInternalsTest
shadowTest project(path: ":beam-runners-core-java", configuration: "shadowTest")
+ shadowTest project(":beam-sdks-java-harness")
shadowTest library.java.avro
shadowTest library.java.kafka_clients
shadowTest library.java.junit
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkJobInvoker.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkJobInvoker.java
new file mode 100644
index 0000000..cea4b07
--- /dev/null
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkJobInvoker.java
@@ -0,0 +1,78 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.spark;
+
+import java.util.UUID;
+import javax.annotation.Nullable;
+import org.apache.beam.model.pipeline.v1.RunnerApi.Pipeline;
+import org.apache.beam.runners.core.construction.PipelineOptionsTranslation;
+import org.apache.beam.runners.fnexecution.jobsubmission.JobInvocation;
+import org.apache.beam.runners.fnexecution.jobsubmission.JobInvoker;
+import org.apache.beam.runners.fnexecution.provisioning.JobInfo;
+import org.apache.beam.vendor.grpc.v1p13p1.com.google.protobuf.Struct;
+import org.apache.beam.vendor.guava.v20_0.com.google.common.util.concurrent.ListeningExecutorService;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** Creates a job invocation to manage the Spark runner's execution of a portable pipeline. */
+public class SparkJobInvoker extends JobInvoker {
+
+ private static final Logger LOG = LoggerFactory.getLogger(SparkJobInvoker.class);
+
+ public static SparkJobInvoker create() {
+ return new SparkJobInvoker();
+ }
+
+ private SparkJobInvoker() {
+ super("spark-runner-job-invoker");
+ }
+
+ @Override
+ protected JobInvocation invokeWithExecutor(
+ Pipeline pipeline,
+ Struct options,
+ @Nullable String retrievalToken,
+ ListeningExecutorService executorService) {
+ LOG.trace("Parsing pipeline options");
+ SparkPipelineOptions sparkOptions =
+ PipelineOptionsTranslation.fromProto(options).as(SparkPipelineOptions.class);
+
+ String invocationId =
+ String.format("%s_%s", sparkOptions.getJobName(), UUID.randomUUID().toString());
+ LOG.info("Invoking job {}", invocationId);
+
+ return createJobInvocation(
+ invocationId, retrievalToken, executorService, pipeline, sparkOptions);
+ }
+
+ static JobInvocation createJobInvocation(
+ String invocationId,
+ String retrievalToken,
+ ListeningExecutorService executorService,
+ Pipeline pipeline,
+ SparkPipelineOptions sparkOptions) {
+ JobInfo jobInfo =
+ JobInfo.create(
+ invocationId,
+ sparkOptions.getJobName(),
+ retrievalToken,
+ PipelineOptionsTranslation.toProto(sparkOptions));
+ SparkPipelineRunner pipelineRunner = new SparkPipelineRunner(sparkOptions);
+ return new JobInvocation(jobInfo, executorService, pipeline, pipelineRunner);
+ }
+}
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkJobServerDriver.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkJobServerDriver.java
new file mode 100644
index 0000000..387907f
--- /dev/null
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkJobServerDriver.java
@@ -0,0 +1,86 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.spark;
+
+import org.apache.beam.runners.fnexecution.ServerFactory;
+import org.apache.beam.runners.fnexecution.jobsubmission.JobInvoker;
+import org.apache.beam.runners.fnexecution.jobsubmission.JobServerDriver;
+import org.apache.beam.sdk.io.FileSystems;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.kohsuke.args4j.CmdLineException;
+import org.kohsuke.args4j.CmdLineParser;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** Driver program that starts a job server for the Spark runner. */
+public class SparkJobServerDriver extends JobServerDriver {
+
+ @Override
+ protected JobInvoker createJobInvoker() {
+ return SparkJobInvoker.create();
+ }
+
+ private static final Logger LOG = LoggerFactory.getLogger(SparkJobServerDriver.class);
+
+ public static void main(String[] args) {
+ FileSystems.setDefaultPipelineOptions(PipelineOptionsFactory.create());
+ fromParams(args).run();
+ }
+
+ private static void printUsage(CmdLineParser parser) {
+ System.err.println(
+ String.format("Usage: java %s arguments...", SparkJobServerDriver.class.getSimpleName()));
+ parser.printUsage(System.err);
+ System.err.println();
+ }
+
+ private static SparkJobServerDriver fromParams(String[] args) {
+ ServerConfiguration configuration = new ServerConfiguration();
+ CmdLineParser parser = new CmdLineParser(configuration);
+ try {
+ parser.parseArgument(args);
+ } catch (CmdLineException e) {
+ LOG.error("Unable to parse command line arguments.", e);
+ printUsage(parser);
+ throw new IllegalArgumentException("Unable to parse command line arguments.", e);
+ }
+
+ return fromConfig(configuration);
+ }
+
+ private static SparkJobServerDriver fromConfig(ServerConfiguration configuration) {
+ return create(
+ configuration,
+ createJobServerFactory(configuration),
+ createArtifactServerFactory(configuration));
+ }
+
+ private static SparkJobServerDriver create(
+ ServerConfiguration configuration,
+ ServerFactory jobServerFactory,
+ ServerFactory artifactServerFactory) {
+ return new SparkJobServerDriver(configuration, jobServerFactory, artifactServerFactory);
+ }
+
+ private SparkJobServerDriver(
+ ServerConfiguration configuration,
+ ServerFactory jobServerFactory,
+ ServerFactory artifactServerFactory) {
+ super(configuration, jobServerFactory, artifactServerFactory);
+ }
+}
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
new file mode 100644
index 0000000..b4d25ab
--- /dev/null
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkPipelineRunner.java
@@ -0,0 +1,79 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.spark;
+
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import org.apache.beam.model.pipeline.v1.RunnerApi;
+import org.apache.beam.model.pipeline.v1.RunnerApi.Pipeline;
+import org.apache.beam.runners.core.construction.graph.GreedyPipelineFuser;
+import org.apache.beam.runners.core.construction.graph.PipelineTrimmer;
+import org.apache.beam.runners.fnexecution.jobsubmission.PortablePipelineRunner;
+import org.apache.beam.runners.fnexecution.provisioning.JobInfo;
+import org.apache.beam.runners.spark.aggregators.AggregatorsAccumulator;
+import org.apache.beam.runners.spark.translation.SparkBatchPortablePipelineTranslator;
+import org.apache.beam.runners.spark.translation.SparkContextFactory;
+import org.apache.beam.runners.spark.translation.SparkTranslationContext;
+import org.apache.spark.api.java.JavaSparkContext;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** Runs a portable pipeline on Apache Spark. */
+public class SparkPipelineRunner implements PortablePipelineRunner {
+
+ private static final Logger LOG = LoggerFactory.getLogger(SparkPipelineRunner.class);
+
+ private final SparkPipelineOptions pipelineOptions;
+
+ public SparkPipelineRunner(SparkPipelineOptions pipelineOptions) {
+ this.pipelineOptions = pipelineOptions;
+ }
+
+ @Override
+ public SparkPipelineResult run(RunnerApi.Pipeline pipeline, JobInfo jobInfo) {
+ SparkBatchPortablePipelineTranslator translator = new SparkBatchPortablePipelineTranslator();
+
+ // Don't let the fuser fuse any subcomponents of native transforms.
+ Pipeline trimmedPipeline = PipelineTrimmer.trim(pipeline, translator.knownUrns());
+ Pipeline fusedPipeline = GreedyPipelineFuser.fuse(trimmedPipeline).toPipeline();
+
+ final JavaSparkContext jsc = SparkContextFactory.getSparkContext(pipelineOptions);
+ LOG.info(String.format("Running job %s on Spark master %s", jobInfo.jobId(), jsc.master()));
+ AggregatorsAccumulator.init(pipelineOptions, jsc);
+ final SparkTranslationContext context =
+ new SparkTranslationContext(jsc, pipelineOptions, jobInfo);
+ final ExecutorService executorService = Executors.newSingleThreadExecutor();
+ final Future<?> submissionFuture =
+ executorService.submit(
+ () -> {
+ translator.translate(fusedPipeline, context);
+ LOG.info(
+ String.format(
+ "Job %s: Pipeline translated successfully. Computing outputs",
+ jobInfo.jobId()));
+ context.computeOutputs();
+ LOG.info(String.format("Job %s finished.", jobInfo.jobId()));
+ });
+
+ SparkPipelineResult result = new SparkPipelineResult.BatchMode(submissionFuture, jsc);
+ result.waitUntilFinish();
+ executorService.shutdown();
+ return result;
+ }
+}
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/BoundedDataset.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/BoundedDataset.java
index 9cd805a..1e620e7 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/BoundedDataset.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/BoundedDataset.java
@@ -117,6 +117,6 @@ public class BoundedDataset<T> implements Dataset {
@Override
public void setName(String name) {
- rdd.setName(name);
+ getRDD().setName(name);
}
}
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
new file mode 100644
index 0000000..c65caa4
--- /dev/null
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkBatchPortablePipelineTranslator.java
@@ -0,0 +1,210 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.spark.translation;
+
+import static org.apache.beam.runners.fnexecution.translation.PipelineTranslatorUtils.createOutputMap;
+import static org.apache.beam.runners.fnexecution.translation.PipelineTranslatorUtils.getWindowingStrategy;
+
+import java.io.IOException;
+import java.util.Collections;
+import java.util.Map;
+import java.util.Set;
+import javax.annotation.Nullable;
+import org.apache.beam.model.pipeline.v1.RunnerApi;
+import org.apache.beam.model.pipeline.v1.RunnerApi.PCollection;
+import org.apache.beam.runners.core.SystemReduceFn;
+import org.apache.beam.runners.core.construction.PTransformTranslation;
+import org.apache.beam.runners.core.construction.graph.ExecutableStage;
+import org.apache.beam.runners.core.construction.graph.PipelineNode;
+import org.apache.beam.runners.core.construction.graph.PipelineNode.PCollectionNode;
+import org.apache.beam.runners.core.construction.graph.PipelineNode.PTransformNode;
+import org.apache.beam.runners.core.construction.graph.QueryablePipeline;
+import org.apache.beam.runners.fnexecution.wire.WireCoders;
+import org.apache.beam.runners.spark.SparkPipelineOptions;
+import org.apache.beam.runners.spark.aggregators.AggregatorsAccumulator;
+import org.apache.beam.sdk.coders.ByteArrayCoder;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.transforms.join.RawUnionValue;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.TimestampCombiner;
+import org.apache.beam.sdk.transforms.windowing.WindowFn;
+import org.apache.beam.sdk.util.WindowedValue;
+import org.apache.beam.sdk.util.WindowedValue.WindowedValueCoder;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.WindowingStrategy;
+import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.BiMap;
+import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.Iterables;
+import org.apache.spark.HashPartitioner;
+import org.apache.spark.Partitioner;
+import org.apache.spark.api.java.JavaRDD;
+
+/** Translates a bounded portable pipeline into a Spark job. */
+public class SparkBatchPortablePipelineTranslator {
+
+ private final ImmutableMap<String, PTransformTranslator> urnToTransformTranslator;
+
+ interface PTransformTranslator {
+
+ /** Translates transformNode from Beam into the Spark context. */
+ void translate(
+ PTransformNode transformNode, RunnerApi.Pipeline pipeline, SparkTranslationContext context);
+ }
+
+ public Set<String> knownUrns() {
+ return urnToTransformTranslator.keySet();
+ }
+
+ public SparkBatchPortablePipelineTranslator() {
+ ImmutableMap.Builder<String, PTransformTranslator> translatorMap = ImmutableMap.builder();
+ translatorMap.put(
+ PTransformTranslation.IMPULSE_TRANSFORM_URN,
+ SparkBatchPortablePipelineTranslator::translateImpulse);
+ translatorMap.put(
+ PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN,
+ SparkBatchPortablePipelineTranslator::translateGroupByKey);
+ translatorMap.put(
+ ExecutableStage.URN, SparkBatchPortablePipelineTranslator::translateExecutableStage);
+ this.urnToTransformTranslator = translatorMap.build();
+ }
+
+ /** Translates pipeline from Beam into the Spark context. */
+ public void translate(final RunnerApi.Pipeline pipeline, SparkTranslationContext context) {
+ QueryablePipeline p =
+ QueryablePipeline.forTransforms(
+ pipeline.getRootTransformIdsList(), pipeline.getComponents());
+ for (PipelineNode.PTransformNode transformNode : p.getTopologicallyOrderedTransforms()) {
+ urnToTransformTranslator
+ .getOrDefault(
+ transformNode.getTransform().getSpec().getUrn(),
+ SparkBatchPortablePipelineTranslator::urnNotFound)
+ .translate(transformNode, pipeline, context);
+ }
+ }
+
+ private static void urnNotFound(
+ PTransformNode transformNode, RunnerApi.Pipeline pipeline, SparkTranslationContext context) {
+ throw new IllegalArgumentException(
+ String.format(
+ "Transform %s has unknown URN %s",
+ transformNode.getId(), transformNode.getTransform().getSpec().getUrn()));
+ }
+
+ private static void translateImpulse(
+ PTransformNode transformNode, RunnerApi.Pipeline pipeline, SparkTranslationContext context) {
+ BoundedDataset<byte[]> output =
+ new BoundedDataset<>(
+ Collections.singletonList(new byte[0]), context.getSparkContext(), ByteArrayCoder.of());
+ context.pushDataset(getOutputId(transformNode), output);
+ }
+
+ private static <K, V> void translateGroupByKey(
+ PTransformNode transformNode, RunnerApi.Pipeline pipeline, SparkTranslationContext context) {
+
+ RunnerApi.Components components = pipeline.getComponents();
+ String inputId = getInputId(transformNode);
+ PCollection inputPCollection = components.getPcollectionsOrThrow(inputId);
+ Dataset inputDataset = context.popDataset(inputId);
+ JavaRDD<WindowedValue<KV<K, V>>> inputRdd = ((BoundedDataset<KV<K, V>>) inputDataset).getRDD();
+ PCollectionNode inputPCollectionNode = PipelineNode.pCollection(inputId, inputPCollection);
+ WindowedValueCoder<KV<K, V>> inputCoder;
+ try {
+ inputCoder =
+ (WindowedValueCoder)
+ WireCoders.instantiateRunnerWireCoder(inputPCollectionNode, components);
+ } catch (IOException e) {
+ throw new RuntimeException(e);
+ }
+ KvCoder<K, V> inputKvCoder = (KvCoder<K, V>) inputCoder.getValueCoder();
+ Coder<K> inputKeyCoder = inputKvCoder.getKeyCoder();
+ Coder<V> inputValueCoder = inputKvCoder.getValueCoder();
+ WindowingStrategy windowingStrategy = getWindowingStrategy(inputId, components);
+ WindowFn<Object, BoundedWindow> windowFn = windowingStrategy.getWindowFn();
+ WindowedValue.WindowedValueCoder<V> wvCoder =
+ WindowedValue.FullWindowedValueCoder.of(inputValueCoder, windowFn.windowCoder());
+
+ JavaRDD<WindowedValue<KV<K, Iterable<V>>>> groupedByKeyAndWindow;
+ if (windowingStrategy.getWindowFn().isNonMerging()
+ && windowingStrategy.getTimestampCombiner() == TimestampCombiner.END_OF_WINDOW) {
+ // we can have a memory sensitive translation for non-merging windows
+ groupedByKeyAndWindow =
+ GroupNonMergingWindowsFunctions.groupByKeyAndWindow(
+ inputRdd, inputKeyCoder, inputValueCoder, windowingStrategy);
+ } else {
+ Partitioner partitioner = getPartitioner(context);
+ JavaRDD<WindowedValue<KV<K, Iterable<WindowedValue<V>>>>> groupedByKeyOnly =
+ GroupCombineFunctions.groupByKeyOnly(inputRdd, inputKeyCoder, wvCoder, partitioner);
+ // for batch, GroupAlsoByWindow uses an in-memory StateInternals.
+ groupedByKeyAndWindow =
+ groupedByKeyOnly.flatMap(
+ new SparkGroupAlsoByWindowViaOutputBufferFn<>(
+ windowingStrategy,
+ new TranslationUtils.InMemoryStateInternalsFactory<>(),
+ SystemReduceFn.buffering(inputValueCoder),
+ context.serializablePipelineOptions,
+ AggregatorsAccumulator.getInstance()));
+ }
+ context.pushDataset(getOutputId(transformNode), new BoundedDataset<>(groupedByKeyAndWindow));
+ }
+
+ private static <InputT, OutputT> void translateExecutableStage(
+ PTransformNode transformNode, RunnerApi.Pipeline pipeline, SparkTranslationContext context) {
+
+ RunnerApi.ExecutableStagePayload stagePayload;
+ try {
+ stagePayload =
+ RunnerApi.ExecutableStagePayload.parseFrom(
+ transformNode.getTransform().getSpec().getPayload());
+ } catch (IOException e) {
+ throw new RuntimeException(e);
+ }
+ String inputPCollectionId = stagePayload.getInput();
+ Dataset inputDataset = context.popDataset(inputPCollectionId);
+ JavaRDD<WindowedValue<InputT>> inputRdd = ((BoundedDataset<InputT>) inputDataset).getRDD();
+ Map<String, String> outputs = transformNode.getTransform().getOutputsMap();
+ BiMap<String, Integer> outputMap = createOutputMap(outputs.values());
+
+ SparkExecutableStageFunction<InputT> function =
+ new SparkExecutableStageFunction<>(stagePayload, context.jobInfo, outputMap);
+ JavaRDD<RawUnionValue> staged = inputRdd.mapPartitions(function);
+
+ for (String outputId : outputs.values()) {
+ JavaRDD<WindowedValue<OutputT>> outputRdd =
+ staged.flatMap(new SparkExecutableStageExtractionFunction<>(outputMap.get(outputId)));
+ context.pushDataset(outputId, new BoundedDataset<>(outputRdd));
+ }
+ }
+
+ @Nullable
+ private static Partitioner getPartitioner(SparkTranslationContext context) {
+ Long bundleSize =
+ context.serializablePipelineOptions.get().as(SparkPipelineOptions.class).getBundleSize();
+ return (bundleSize > 0)
+ ? null
+ : new HashPartitioner(context.getSparkContext().defaultParallelism());
+ }
+
+ private static String getInputId(PTransformNode transformNode) {
+ return Iterables.getOnlyElement(transformNode.getTransform().getInputsMap().values());
+ }
+
+ private static String getOutputId(PTransformNode transformNode) {
+ return Iterables.getOnlyElement(transformNode.getTransform().getOutputsMap().values());
+ }
+}
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkExecutableStageExtractionFunction.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkExecutableStageExtractionFunction.java
new file mode 100644
index 0000000..a5d1bbd
--- /dev/null
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkExecutableStageExtractionFunction.java
@@ -0,0 +1,42 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.spark.translation;
+
+import java.util.Collections;
+import java.util.Iterator;
+import org.apache.beam.sdk.transforms.join.RawUnionValue;
+import org.apache.beam.sdk.util.WindowedValue;
+import org.apache.spark.api.java.function.FlatMapFunction;
+
+class SparkExecutableStageExtractionFunction<OutputT>
+ implements FlatMapFunction<RawUnionValue, WindowedValue<OutputT>> {
+ private final int unionTag;
+
+ SparkExecutableStageExtractionFunction(int unionTag) {
+ this.unionTag = unionTag;
+ }
+
+ @Override
+ public Iterator<WindowedValue<OutputT>> call(RawUnionValue rawUnionValue) {
+ if (rawUnionValue.getUnionTag() == unionTag) {
+ WindowedValue<OutputT> output = (WindowedValue<OutputT>) rawUnionValue.getValue();
+ return Collections.singleton(output).iterator();
+ }
+ return Collections.emptyIterator();
+ }
+}
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
new file mode 100644
index 0000000..93250bc
--- /dev/null
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkExecutableStageFunction.java
@@ -0,0 +1,155 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.spark.translation;
+
+import java.io.Serializable;
+import java.util.EnumMap;
+import java.util.Iterator;
+import java.util.Locale;
+import java.util.Map;
+import java.util.concurrent.ConcurrentLinkedQueue;
+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;
+import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateKey.TypeCase;
+import org.apache.beam.model.pipeline.v1.RunnerApi;
+import org.apache.beam.runners.core.construction.graph.ExecutableStage;
+import org.apache.beam.runners.fnexecution.control.BundleProgressHandler;
+import org.apache.beam.runners.fnexecution.control.DefaultJobBundleFactory;
+import org.apache.beam.runners.fnexecution.control.JobBundleFactory;
+import org.apache.beam.runners.fnexecution.control.OutputReceiverFactory;
+import org.apache.beam.runners.fnexecution.control.RemoteBundle;
+import org.apache.beam.runners.fnexecution.control.StageBundleFactory;
+import org.apache.beam.runners.fnexecution.provisioning.JobInfo;
+import org.apache.beam.runners.fnexecution.state.StateRequestHandler;
+import org.apache.beam.runners.fnexecution.state.StateRequestHandlers;
+import org.apache.beam.sdk.fn.data.FnDataReceiver;
+import org.apache.beam.sdk.transforms.join.RawUnionValue;
+import org.apache.beam.sdk.util.WindowedValue;
+import org.apache.spark.api.java.function.FlatMapFunction;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Spark function that passes its input through an SDK-executed {@link
+ * org.apache.beam.runners.core.construction.graph.ExecutableStage}.
+ *
+ * <p>The output of this operation is a multiplexed {@link Dataset} whose elements are tagged with a
+ * union coder. The coder's tags are determined by {@link SparkExecutableStageFunction#outputMap}.
+ * The resulting data set should be further processed by a {@link
+ * SparkExecutableStageExtractionFunction}.
+ */
+public class SparkExecutableStageFunction<InputT>
+ implements FlatMapFunction<Iterator<WindowedValue<InputT>>, RawUnionValue> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(SparkExecutableStageFunction.class);
+
+ private final RunnerApi.ExecutableStagePayload stagePayload;
+ private final Map<String, Integer> outputMap;
+ private final JobBundleFactoryCreator jobBundleFactoryCreator;
+
+ SparkExecutableStageFunction(
+ RunnerApi.ExecutableStagePayload stagePayload,
+ JobInfo jobInfo,
+ Map<String, Integer> outputMap) {
+ this(stagePayload, outputMap, () -> DefaultJobBundleFactory.create(jobInfo));
+ }
+
+ SparkExecutableStageFunction(
+ RunnerApi.ExecutableStagePayload stagePayload,
+ Map<String, Integer> outputMap,
+ JobBundleFactoryCreator jobBundleFactoryCreator) {
+ this.stagePayload = stagePayload;
+ this.outputMap = outputMap;
+ this.jobBundleFactoryCreator = jobBundleFactoryCreator;
+ }
+
+ @Override
+ public Iterator<RawUnionValue> call(Iterator<WindowedValue<InputT>> inputs) throws Exception {
+ JobBundleFactory jobBundleFactory = jobBundleFactoryCreator.create();
+ ExecutableStage executableStage = ExecutableStage.fromPayload(stagePayload);
+ try (StageBundleFactory stageBundleFactory = jobBundleFactory.forStage(executableStage)) {
+ ConcurrentLinkedQueue<RawUnionValue> collector = new ConcurrentLinkedQueue<>();
+ ReceiverFactory receiverFactory = new ReceiverFactory(collector, outputMap);
+ EnumMap<TypeCase, StateRequestHandler> handlers = new EnumMap<>(StateKey.TypeCase.class);
+ // TODO add state request handlers
+ StateRequestHandler stateRequestHandler =
+ StateRequestHandlers.delegateBasedUponType(handlers);
+ SparkBundleProgressHandler bundleProgressHandler = new SparkBundleProgressHandler();
+ try (RemoteBundle bundle =
+ stageBundleFactory.getBundle(
+ receiverFactory, stateRequestHandler, bundleProgressHandler)) {
+ String inputPCollectionId = executableStage.getInputPCollection().getId();
+ FnDataReceiver<WindowedValue<?>> mainReceiver =
+ bundle.getInputReceivers().get(inputPCollectionId);
+ while (inputs.hasNext()) {
+ WindowedValue<InputT> input = inputs.next();
+ mainReceiver.accept(input);
+ }
+ }
+ return collector.iterator();
+ } catch (Exception e) {
+ LOG.error("Spark executable stage fn terminated with exception: ", e);
+ throw e;
+ }
+ }
+
+ interface JobBundleFactoryCreator extends Serializable {
+ JobBundleFactory create();
+ }
+
+ /**
+ * Receiver factory that wraps outgoing elements with the corresponding union tag for a
+ * multiplexed PCollection.
+ */
+ private static class ReceiverFactory implements OutputReceiverFactory {
+
+ private final ConcurrentLinkedQueue<RawUnionValue> collector;
+ private final Map<String, Integer> outputMap;
+
+ ReceiverFactory(
+ ConcurrentLinkedQueue<RawUnionValue> collector, Map<String, Integer> outputMap) {
+ this.collector = collector;
+ this.outputMap = outputMap;
+ }
+
+ @Override
+ public <OutputT> FnDataReceiver<OutputT> create(String pCollectionId) {
+ Integer unionTag = outputMap.get(pCollectionId);
+ if (unionTag == null) {
+ throw new IllegalStateException(
+ String.format(Locale.ENGLISH, "Unknown PCollectionId %s", pCollectionId));
+ }
+ int tagInt = unionTag;
+ return receivedElement -> collector.add(new RawUnionValue(tagInt, receivedElement));
+ }
+ }
+
+ private static class SparkBundleProgressHandler implements BundleProgressHandler {
+
+ @Override
+ public void onProgress(ProcessBundleProgressResponse progress) {
+ // TODO
+ }
+
+ @Override
+ public void onCompleted(ProcessBundleResponse response) {
+ // TODO
+ }
+ }
+}
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkTranslationContext.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkTranslationContext.java
new file mode 100644
index 0000000..237e735
--- /dev/null
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkTranslationContext.java
@@ -0,0 +1,71 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.spark.translation;
+
+import java.util.LinkedHashMap;
+import java.util.LinkedHashSet;
+import java.util.Map;
+import java.util.Set;
+import org.apache.beam.runners.core.construction.SerializablePipelineOptions;
+import org.apache.beam.runners.fnexecution.provisioning.JobInfo;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.spark.api.java.JavaSparkContext;
+
+/**
+ * Translation context used to lazily store Spark data sets during portable pipeline translation and
+ * compute them after translation.
+ */
+public class SparkTranslationContext {
+ private final JavaSparkContext jsc;
+ final JobInfo jobInfo;
+ private final Map<String, Dataset> datasets = new LinkedHashMap<>();
+ private final Set<Dataset> leaves = new LinkedHashSet<>();
+ final SerializablePipelineOptions serializablePipelineOptions;
+
+ public SparkTranslationContext(JavaSparkContext jsc, PipelineOptions options, JobInfo jobInfo) {
+ this.jsc = jsc;
+ this.serializablePipelineOptions = new SerializablePipelineOptions(options);
+ this.jobInfo = jobInfo;
+ }
+
+ public JavaSparkContext getSparkContext() {
+ return jsc;
+ }
+
+ /** Add output of transform to context. */
+ public void pushDataset(String pCollectionId, Dataset dataset) {
+ dataset.setName(pCollectionId);
+ // TODO cache
+ datasets.put(pCollectionId, dataset);
+ leaves.add(dataset);
+ }
+
+ /** Retrieve the dataset for the pCollection id and remove it from the DAG's leaves. */
+ public Dataset popDataset(String pCollectionId) {
+ Dataset dataset = datasets.get(pCollectionId);
+ leaves.remove(dataset);
+ return dataset;
+ }
+
+ /** Compute the outputs for all RDDs that are leaves in the DAG. */
+ public void computeOutputs() {
+ for (Dataset dataset : leaves) {
+ dataset.action(); // force computation.
+ }
+ }
+}
diff --git a/runners/spark/src/test/java/org/apache/beam/runners/spark/SparkPortableExecutionTest.java b/runners/spark/src/test/java/org/apache/beam/runners/spark/SparkPortableExecutionTest.java
new file mode 100644
index 0000000..ad97ec0
--- /dev/null
+++ b/runners/spark/src/test/java/org/apache/beam/runners/spark/SparkPortableExecutionTest.java
@@ -0,0 +1,139 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.spark;
+
+import java.io.Serializable;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import org.apache.beam.model.jobmanagement.v1.JobApi.JobState.Enum;
+import org.apache.beam.model.pipeline.v1.RunnerApi;
+import org.apache.beam.runners.core.construction.Environments;
+import org.apache.beam.runners.core.construction.PipelineTranslation;
+import org.apache.beam.runners.fnexecution.jobsubmission.JobInvocation;
+import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.options.PortablePipelineOptions;
+import org.apache.beam.sdk.testing.CrashingRunner;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.GroupByKey;
+import org.apache.beam.sdk.transforms.Impulse;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.WithKeys;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.vendor.guava.v20_0.com.google.common.util.concurrent.ListeningExecutorService;
+import org.apache.beam.vendor.guava.v20_0.com.google.common.util.concurrent.MoreExecutors;
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Tests the execution of a pipeline from specification to execution on the portable Spark runner.
+ */
+public class SparkPortableExecutionTest implements Serializable {
+
+ private static final Logger LOG = LoggerFactory.getLogger(SparkPortableExecutionTest.class);
+
+ private static ListeningExecutorService sparkJobExecutor;
+
+ @BeforeClass
+ public static void setup() {
+ // Restrict this to only one thread to avoid multiple Spark clusters up at the same time
+ // which is not suitable for memory-constraint environments, i.e. Jenkins.
+ sparkJobExecutor = MoreExecutors.listeningDecorator(Executors.newFixedThreadPool(1));
+ }
+
+ @AfterClass
+ public static void tearDown() throws InterruptedException {
+ sparkJobExecutor.shutdown();
+ sparkJobExecutor.awaitTermination(10, TimeUnit.SECONDS);
+ if (!sparkJobExecutor.isShutdown()) {
+ LOG.warn("Could not shut down Spark job executor");
+ }
+ sparkJobExecutor = null;
+ }
+
+ @Test(timeout = 120_000)
+ public void testExecution() throws Exception {
+ PipelineOptions options = PipelineOptionsFactory.create();
+ options.setRunner(CrashingRunner.class);
+ options
+ .as(PortablePipelineOptions.class)
+ .setDefaultEnvironmentType(Environments.ENVIRONMENT_EMBEDDED);
+
+ Pipeline p = Pipeline.create(options);
+ PCollection<KV<String, Iterable<Long>>> result =
+ p.apply("impulse", Impulse.create())
+ .apply(
+ "create",
+ ParDo.of(
+ new DoFn<byte[], String>() {
+ @ProcessElement
+ public void process(ProcessContext context) {
+ context.output("zero");
+ context.output("one");
+ context.output("two");
+ }
+ }))
+ .apply(
+ "len",
+ ParDo.of(
+ new DoFn<String, Long>() {
+ @ProcessElement
+ public void process(ProcessContext context) {
+ context.output((long) context.element().length());
+ }
+ }))
+ .apply("addKeys", WithKeys.of("foo"))
+ // First GBK just to verify GBK works
+ .apply("gbk", GroupByKey.create())
+ .apply(
+ "print",
+ ParDo.of(
+ new DoFn<KV<String, Iterable<Long>>, KV<String, Long>>() {
+ @ProcessElement
+ public void process(ProcessContext context) {
+ LOG.info("Output element: {}", context.element());
+ for (Long i : context.element().getValue()) {
+ context.output(KV.of(context.element().getKey(), i));
+ }
+ }
+ }))
+ // Second GBK forces the output to be materialized
+ .apply("gbk", GroupByKey.create());
+
+ // TODO Get PAssert working to test pipeline result
+
+ RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(p);
+
+ JobInvocation jobInvocation =
+ SparkJobInvoker.createJobInvocation(
+ "fakeId",
+ "fakeRetrievalToken",
+ sparkJobExecutor,
+ pipelineProto,
+ options.as(SparkPipelineOptions.class));
+ jobInvocation.start();
+ while (jobInvocation.getState() != Enum.DONE) {
+ Thread.sleep(1000);
+ }
+ }
+}
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
new file mode 100644
index 0000000..8f1bdca
--- /dev/null
+++ b/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/SparkExecutableStageFunctionTest.java
@@ -0,0 +1,203 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.spark.translation;
+
+import static org.apache.beam.runners.core.construction.PTransformTranslation.PAR_DO_TRANSFORM_URN;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.contains;
+import static org.mockito.Matchers.any;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+import static org.mockito.Mockito.when;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.Iterator;
+import java.util.Map;
+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;
+import org.apache.beam.model.pipeline.v1.RunnerApi.PCollection;
+import org.apache.beam.runners.fnexecution.control.BundleProgressHandler;
+import org.apache.beam.runners.fnexecution.control.JobBundleFactory;
+import org.apache.beam.runners.fnexecution.control.OutputReceiverFactory;
+import org.apache.beam.runners.fnexecution.control.ProcessBundleDescriptors;
+import org.apache.beam.runners.fnexecution.control.RemoteBundle;
+import org.apache.beam.runners.fnexecution.control.StageBundleFactory;
+import org.apache.beam.runners.fnexecution.state.StateRequestHandler;
+import org.apache.beam.runners.spark.translation.SparkExecutableStageFunction.JobBundleFactoryCreator;
+import org.apache.beam.sdk.fn.data.FnDataReceiver;
+import org.apache.beam.sdk.transforms.join.RawUnionValue;
+import org.apache.beam.sdk.util.WindowedValue;
+import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ImmutableMap;
+import org.junit.Before;
+import org.junit.Test;
+import org.mockito.Mock;
+import org.mockito.Mockito;
+import org.mockito.MockitoAnnotations;
+
+/** Unit tests for {@link SparkExecutableStageFunction}. */
+public class SparkExecutableStageFunctionTest {
+ @Mock private JobBundleFactoryCreator jobBundleFactoryCreator;
+ @Mock private JobBundleFactory jobBundleFactory;
+ @Mock private StageBundleFactory stageBundleFactory;
+ @Mock private RemoteBundle remoteBundle;
+
+ private final String inputId = "input-id";
+ private final ExecutableStagePayload stagePayload =
+ ExecutableStagePayload.newBuilder()
+ .setInput(inputId)
+ .setComponents(
+ Components.newBuilder()
+ .putTransforms(
+ "transform-id",
+ RunnerApi.PTransform.newBuilder()
+ .putInputs("input-name", inputId)
+ .setSpec(RunnerApi.FunctionSpec.newBuilder().setUrn(PAR_DO_TRANSFORM_URN))
+ .build())
+ .putPcollections(inputId, PCollection.getDefaultInstance())
+ .build())
+ .build();
+
+ @Before
+ public void setUpMocks() throws Exception {
+ MockitoAnnotations.initMocks(this);
+ when(jobBundleFactoryCreator.create()).thenReturn(jobBundleFactory);
+ when(jobBundleFactory.forStage(any())).thenReturn(stageBundleFactory);
+ when(stageBundleFactory.getBundle(any(), any(), any())).thenReturn(remoteBundle);
+ @SuppressWarnings("unchecked")
+ ImmutableMap<String, FnDataReceiver<WindowedValue<?>>> inputReceiver =
+ ImmutableMap.of("input", Mockito.mock(FnDataReceiver.class));
+ when(remoteBundle.getInputReceivers()).thenReturn(inputReceiver);
+ }
+
+ @Test(expected = Exception.class)
+ public void sdkErrorsSurfaceOnClose() throws Exception {
+ SparkExecutableStageFunction<Integer> function = getFunction(Collections.emptyMap());
+ doThrow(new Exception()).when(remoteBundle).close();
+ function.call(Collections.emptyIterator());
+ }
+
+ @Test
+ public void expectedInputsAreSent() throws Exception {
+ SparkExecutableStageFunction<Integer> function = getFunction(Collections.emptyMap());
+
+ RemoteBundle bundle = Mockito.mock(RemoteBundle.class);
+ when(stageBundleFactory.getBundle(any(), any(), any())).thenReturn(bundle);
+
+ @SuppressWarnings("unchecked")
+ FnDataReceiver<WindowedValue<?>> receiver = Mockito.mock(FnDataReceiver.class);
+ when(bundle.getInputReceivers()).thenReturn(ImmutableMap.of(inputId, receiver));
+
+ WindowedValue<Integer> one = WindowedValue.valueInGlobalWindow(1);
+ WindowedValue<Integer> two = WindowedValue.valueInGlobalWindow(2);
+ WindowedValue<Integer> three = WindowedValue.valueInGlobalWindow(3);
+ function.call(Arrays.asList(one, two, three).iterator());
+
+ verify(receiver).accept(one);
+ verify(receiver).accept(two);
+ verify(receiver).accept(three);
+ verifyNoMoreInteractions(receiver);
+ }
+
+ @Test
+ public void outputsAreTaggedCorrectly() throws Exception {
+ WindowedValue<Integer> three = WindowedValue.valueInGlobalWindow(3);
+ WindowedValue<Integer> four = WindowedValue.valueInGlobalWindow(4);
+ WindowedValue<Integer> five = WindowedValue.valueInGlobalWindow(5);
+ Map<String, Integer> outputTagMap =
+ ImmutableMap.of(
+ "one", 1,
+ "two", 2,
+ "three", 3);
+
+ // We use a real StageBundleFactory here in order to exercise the output receiver factory.
+ StageBundleFactory stageBundleFactory =
+ new StageBundleFactory() {
+
+ private boolean once;
+
+ @Override
+ public RemoteBundle getBundle(
+ OutputReceiverFactory receiverFactory,
+ StateRequestHandler stateRequestHandler,
+ BundleProgressHandler progressHandler) {
+ return new RemoteBundle() {
+ @Override
+ public String getId() {
+ return "bundle-id";
+ }
+
+ @Override
+ public Map<String, FnDataReceiver<WindowedValue<?>>> getInputReceivers() {
+ return ImmutableMap.of(
+ "input",
+ input -> {
+ /* Ignore input*/
+ });
+ }
+
+ @Override
+ public void close() throws Exception {
+ if (once) {
+ return;
+ }
+ // Emit all values to the runner when the bundle is closed.
+ receiverFactory.create("one").accept(three);
+ receiverFactory.create("two").accept(four);
+ receiverFactory.create("three").accept(five);
+ once = true;
+ }
+ };
+ }
+
+ @Override
+ public ProcessBundleDescriptors.ExecutableProcessBundleDescriptor
+ getProcessBundleDescriptor() {
+ return Mockito.mock(ProcessBundleDescriptors.ExecutableProcessBundleDescriptor.class);
+ }
+
+ @Override
+ public void close() {}
+ };
+ when(jobBundleFactory.forStage(any())).thenReturn(stageBundleFactory);
+
+ SparkExecutableStageFunction<Integer> function = getFunction(outputTagMap);
+ Iterator<RawUnionValue> iterator = function.call(Collections.emptyIterator());
+ Iterable<RawUnionValue> iterable = () -> iterator;
+
+ assertThat(
+ iterable,
+ contains(
+ new RawUnionValue(1, three), new RawUnionValue(2, four), new RawUnionValue(3, five)));
+ }
+
+ @Test
+ public void testStageBundleClosed() throws Exception {
+ SparkExecutableStageFunction<Integer> function = getFunction(Collections.emptyMap());
+ function.call(Collections.emptyIterator());
+ verify(stageBundleFactory).getBundle(any(), any(), any());
+ verify(stageBundleFactory).close();
+ verifyNoMoreInteractions(stageBundleFactory);
+ }
+
+ private <T> SparkExecutableStageFunction<T> getFunction(Map<String, Integer> outputMap) {
+ return new SparkExecutableStageFunction<>(stagePayload, outputMap, jobBundleFactoryCreator);
+ }
+}