| /* |
| * 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.List; |
| 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.ExecutableStagePayload.SideInputId; |
| 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.ReadTranslation; |
| 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.runners.spark.io.SourceRDD; |
| import org.apache.beam.runners.spark.metrics.MetricsAccumulator; |
| 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.io.BoundedSource; |
| 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.ImmutableSet; |
| import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.Iterables; |
| import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.Sets; |
| import org.apache.spark.HashPartitioner; |
| import org.apache.spark.Partitioner; |
| import org.apache.spark.api.java.JavaRDD; |
| import org.apache.spark.api.java.JavaSparkContext; |
| import org.apache.spark.broadcast.Broadcast; |
| import scala.Tuple2; |
| |
| /** 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() { |
| // Do not expose Read as a known URN because we only want to support Read |
| // through the Java ExpansionService. We can't translate Reads for other |
| // languages. |
| return Sets.difference( |
| urnToTransformTranslator.keySet(), |
| ImmutableSet.of(PTransformTranslation.READ_TRANSFORM_URN)); |
| } |
| |
| 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); |
| translatorMap.put( |
| PTransformTranslation.FLATTEN_TRANSFORM_URN, |
| SparkBatchPortablePipelineTranslator::translateFlatten); |
| translatorMap.put( |
| PTransformTranslation.READ_TRANSFORM_URN, |
| SparkBatchPortablePipelineTranslator::translateRead); |
| 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, SideInputT> 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()); |
| |
| ImmutableMap.Builder<String, Tuple2<Broadcast<List<byte[]>>, WindowedValueCoder<SideInputT>>> |
| broadcastVariablesBuilder = ImmutableMap.builder(); |
| for (SideInputId sideInputId : stagePayload.getSideInputsList()) { |
| RunnerApi.Components components = stagePayload.getComponents(); |
| String collectionId = |
| components |
| .getTransformsOrThrow(sideInputId.getTransformId()) |
| .getInputsOrThrow(sideInputId.getLocalName()); |
| Tuple2<Broadcast<List<byte[]>>, WindowedValueCoder<SideInputT>> tuple2 = |
| broadcastSideInput(collectionId, components, context); |
| broadcastVariablesBuilder.put(collectionId, tuple2); |
| } |
| |
| SparkExecutableStageFunction<InputT, SideInputT> function = |
| new SparkExecutableStageFunction<>( |
| stagePayload, |
| context.jobInfo, |
| outputMap, |
| broadcastVariablesBuilder.build(), |
| MetricsAccumulator.getInstance()); |
| 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)); |
| } |
| if (outputs.isEmpty()) { |
| // After pipeline translation, we traverse the set of unconsumed PCollections and add a |
| // no-op sink to each to make sure they are materialized by Spark. However, some SDK-executed |
| // stages have no runner-visible output after fusion. We handle this case by adding a sink |
| // here. |
| JavaRDD<WindowedValue<OutputT>> outputRdd = |
| staged.flatMap((rawUnionValue) -> Collections.emptyIterator()); |
| context.pushDataset( |
| String.format("EmptyOutputSink_%d", context.nextSinkId()), |
| new BoundedDataset<>(outputRdd)); |
| } |
| } |
| |
| /** |
| * Collect and serialize the data and then broadcast the result. *This can be expensive.* |
| * |
| * @return Spark broadcast variable and coder to decode its contents |
| */ |
| private static <T> Tuple2<Broadcast<List<byte[]>>, WindowedValueCoder<T>> broadcastSideInput( |
| String collectionId, RunnerApi.Components components, SparkTranslationContext context) { |
| PCollection collection = components.getPcollectionsOrThrow(collectionId); |
| @SuppressWarnings("unchecked") |
| BoundedDataset<T> dataset = (BoundedDataset<T>) context.popDataset(collectionId); |
| PCollectionNode collectionNode = PipelineNode.pCollection(collectionId, collection); |
| WindowedValueCoder<T> coder; |
| try { |
| coder = |
| (WindowedValueCoder<T>) WireCoders.instantiateRunnerWireCoder(collectionNode, components); |
| } catch (IOException e) { |
| throw new RuntimeException(e); |
| } |
| List<byte[]> bytes = dataset.getBytes(coder); |
| Broadcast<List<byte[]>> broadcast = context.getSparkContext().broadcast(bytes); |
| return new Tuple2<>(broadcast, coder); |
| } |
| |
| private static <T> void translateFlatten( |
| PTransformNode transformNode, RunnerApi.Pipeline pipeline, SparkTranslationContext context) { |
| |
| Map<String, String> inputsMap = transformNode.getTransform().getInputsMap(); |
| |
| JavaRDD<WindowedValue<T>> unionRDD; |
| if (inputsMap.isEmpty()) { |
| unionRDD = context.getSparkContext().emptyRDD(); |
| } else { |
| JavaRDD<WindowedValue<T>>[] rdds = new JavaRDD[inputsMap.size()]; |
| int index = 0; |
| for (String inputId : inputsMap.values()) { |
| rdds[index] = ((BoundedDataset<T>) context.popDataset(inputId)).getRDD(); |
| index++; |
| } |
| unionRDD = context.getSparkContext().union(rdds); |
| } |
| context.pushDataset(getOutputId(transformNode), new BoundedDataset<>(unionRDD)); |
| } |
| |
| private static <T> void translateRead( |
| PTransformNode transformNode, RunnerApi.Pipeline pipeline, SparkTranslationContext context) { |
| String stepName = transformNode.getTransform().getUniqueName(); |
| final JavaSparkContext jsc = context.getSparkContext(); |
| |
| BoundedSource boundedSource; |
| try { |
| boundedSource = |
| ReadTranslation.boundedSourceFromProto( |
| RunnerApi.ReadPayload.parseFrom(transformNode.getTransform().getSpec().getPayload())); |
| } catch (IOException e) { |
| throw new RuntimeException("Failed to extract BoundedSource from ReadPayload.", e); |
| } |
| |
| // create an RDD from a BoundedSource. |
| JavaRDD<WindowedValue<T>> input = |
| new SourceRDD.Bounded<>( |
| jsc.sc(), boundedSource, context.serializablePipelineOptions, stepName) |
| .toJavaRDD(); |
| |
| context.pushDataset(getOutputId(transformNode), new BoundedDataset<>(input)); |
| } |
| |
| @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()); |
| } |
| } |
