| /* |
| * 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.nemo.compiler.frontend.beam; |
| |
| import com.google.common.collect.Iterables; |
| import org.apache.beam.runners.core.SystemReduceFn; |
| import org.apache.beam.runners.core.construction.ParDoTranslation; |
| import org.apache.beam.runners.core.construction.TransformInputs; |
| import org.apache.beam.sdk.Pipeline; |
| import org.apache.beam.sdk.coders.CannotProvideCoderException; |
| import org.apache.beam.sdk.coders.Coder; |
| import org.apache.beam.sdk.coders.KvCoder; |
| import org.apache.beam.sdk.io.Read; |
| import org.apache.beam.sdk.runners.AppliedPTransform; |
| import org.apache.beam.sdk.runners.TransformHierarchy; |
| import org.apache.beam.sdk.transforms.*; |
| import org.apache.beam.sdk.transforms.display.DisplayData; |
| import org.apache.beam.sdk.transforms.display.HasDisplayData; |
| import org.apache.beam.sdk.transforms.windowing.GlobalWindows; |
| import org.apache.beam.sdk.transforms.windowing.Window; |
| import org.apache.beam.sdk.transforms.windowing.WindowFn; |
| import org.apache.beam.sdk.util.AppliedCombineFn; |
| import org.apache.beam.sdk.values.PCollection; |
| import org.apache.beam.sdk.values.PCollectionView; |
| import org.apache.beam.sdk.values.TupleTag; |
| import org.apache.beam.sdk.values.TupleTagList; |
| import org.apache.nemo.common.ir.edge.IREdge; |
| import org.apache.nemo.common.ir.edge.executionproperty.CommunicationPatternProperty; |
| import org.apache.nemo.common.ir.vertex.IRVertex; |
| import org.apache.nemo.common.ir.vertex.OperatorVertex; |
| import org.apache.nemo.common.ir.vertex.transform.Transform; |
| import org.apache.nemo.compiler.frontend.beam.source.BeamBoundedSourceVertex; |
| import org.apache.nemo.compiler.frontend.beam.source.BeamUnboundedSourceVertex; |
| import org.apache.nemo.compiler.frontend.beam.transform.*; |
| import org.slf4j.Logger; |
| import org.slf4j.LoggerFactory; |
| |
| import java.io.IOException; |
| import java.lang.annotation.ElementType; |
| import java.lang.annotation.Retention; |
| import java.lang.annotation.RetentionPolicy; |
| import java.lang.annotation.Target; |
| import java.lang.reflect.InvocationTargetException; |
| import java.lang.reflect.Method; |
| import java.util.*; |
| import java.util.function.Function; |
| import java.util.stream.Collectors; |
| import java.util.stream.IntStream; |
| |
| /** |
| * A collection of translators for the Beam PTransforms. |
| */ |
| final class PipelineTranslator { |
| public static final PipelineTranslator INSTANCE = new PipelineTranslator(); |
| private static final Logger LOG = LoggerFactory.getLogger(PipelineTranslator.class.getName()); |
| |
| private final Map<Class<? extends PTransform>, Method> primitiveTransformToTranslator = new HashMap<>(); |
| private final Map<Class<? extends PTransform>, Method> compositeTransformToTranslator = new HashMap<>(); |
| |
| /** |
| * Creates the translator, while building a map between {@link PTransform}s and the corresponding translators. |
| */ |
| private PipelineTranslator() { |
| for (final Method translator : getClass().getDeclaredMethods()) { |
| final PrimitiveTransformTranslator primitive = translator.getAnnotation(PrimitiveTransformTranslator.class); |
| final CompositeTransformTranslator composite = translator.getAnnotation(CompositeTransformTranslator.class); |
| if (primitive != null) { |
| for (final Class<? extends PTransform> transform : primitive.value()) { |
| if (primitiveTransformToTranslator.containsKey(transform)) { |
| throw new RuntimeException(String.format("Translator for primitive transform %s is" |
| + "already registered: %s", transform, primitiveTransformToTranslator.get(transform))); |
| } |
| primitiveTransformToTranslator.put(transform, translator); |
| } |
| } |
| if (composite != null) { |
| for (final Class<? extends PTransform> transform : composite.value()) { |
| if (compositeTransformToTranslator.containsKey(transform)) { |
| throw new RuntimeException(String.format("Translator for composite transform %s is" |
| + "already registered: %s", transform, compositeTransformToTranslator.get(transform))); |
| } |
| compositeTransformToTranslator.put(transform, translator); |
| } |
| } |
| } |
| } |
| |
| /** |
| * @param context provides translation context. |
| * @param primitive primitive node. |
| */ |
| void translatePrimitive(final PipelineTranslationContext context, |
| final TransformHierarchy.Node primitive) { |
| final PTransform<?, ?> transform = primitive.getTransform(); |
| Class<?> clazz = transform.getClass(); |
| final Method translator = primitiveTransformToTranslator.get(clazz); |
| if (translator == null) { |
| throw new UnsupportedOperationException( |
| String.format("Primitive transform %s is not supported", transform.getClass().getCanonicalName())); |
| } else { |
| try { |
| translator.setAccessible(true); |
| translator.invoke(null, context, primitive, transform); |
| } catch (final IllegalAccessException e) { |
| throw new RuntimeException(e); |
| } catch (final InvocationTargetException | RuntimeException e) { |
| throw new RuntimeException(String.format( |
| "Translator %s have failed to translate %s", translator, transform), e); |
| } |
| } |
| } |
| |
| /** |
| * @param context context. |
| * @param composite transform. |
| * @return behavior controls whether or not child transforms are visited. |
| */ |
| Pipeline.PipelineVisitor.CompositeBehavior translateComposite(final PipelineTranslationContext context, |
| final TransformHierarchy.Node composite) { |
| final PTransform<?, ?> transform = composite.getTransform(); |
| if (transform == null) { |
| // root beam node |
| return Pipeline.PipelineVisitor.CompositeBehavior.ENTER_TRANSFORM; |
| } |
| |
| Class<?> clazz = transform.getClass(); |
| final Method translator = compositeTransformToTranslator.get(clazz); |
| if (translator == null) { |
| return Pipeline.PipelineVisitor.CompositeBehavior.ENTER_TRANSFORM; |
| } else { |
| try { |
| translator.setAccessible(true); |
| return (Pipeline.PipelineVisitor.CompositeBehavior) translator.invoke(null, context, composite, transform); |
| } catch (final IllegalAccessException e) { |
| throw new RuntimeException(e); |
| } catch (final InvocationTargetException | RuntimeException e) { |
| throw new RuntimeException(String.format( |
| "Translator %s have failed to translate %s", translator, transform), e); |
| } |
| } |
| } |
| |
| /** |
| * Annotates translator for PrimitiveTransform. |
| */ |
| @Target(ElementType.METHOD) |
| @Retention(RetentionPolicy.RUNTIME) |
| private @interface PrimitiveTransformTranslator { |
| /** |
| * @return primitive transform. |
| */ |
| Class<? extends PTransform>[] value(); |
| } |
| |
| /** |
| * Annotates translator for CompositeTransform. |
| */ |
| @Target(ElementType.METHOD) |
| @Retention(RetentionPolicy.RUNTIME) |
| private @interface CompositeTransformTranslator { |
| /** |
| * @return composite transform. |
| */ |
| Class<? extends PTransform>[] value(); |
| } |
| |
| //////////////////////////////////////////////////////////////////////////////////////////////////////// |
| /////////////////////// PRIMITIVE TRANSFORMS |
| |
| /** |
| * @param ctx provides translation context |
| * @param beamNode the beam node to be translated |
| * @param transform transform which can be obtained from {@code beamNode} |
| */ |
| @PrimitiveTransformTranslator(Read.Unbounded.class) |
| private static void unboundedReadTranslator(final PipelineTranslationContext ctx, |
| final TransformHierarchy.Node beamNode, |
| final Read.Unbounded<?> transform) { |
| final IRVertex vertex = new BeamUnboundedSourceVertex<>(transform.getSource(), DisplayData.from(transform)); |
| ctx.addVertex(vertex); |
| beamNode.getInputs().values().forEach(input -> ctx.addEdgeTo(vertex, input)); |
| beamNode.getOutputs().values().forEach(output -> ctx.registerMainOutputFrom(beamNode, vertex, output)); |
| } |
| |
| /** |
| * @param ctx provides translation context |
| * @param beamNode the beam node to be translated |
| * @param transform transform which can be obtained from {@code beamNode} |
| */ |
| @PrimitiveTransformTranslator(Read.Bounded.class) |
| private static void boundedReadTranslator(final PipelineTranslationContext ctx, |
| final TransformHierarchy.Node beamNode, |
| final Read.Bounded<?> transform) { |
| final IRVertex vertex = new BeamBoundedSourceVertex<>(transform.getSource(), DisplayData.from(transform)); |
| ctx.addVertex(vertex); |
| beamNode.getInputs().values().forEach(input -> ctx.addEdgeTo(vertex, input)); |
| beamNode.getOutputs().values().forEach(output -> ctx.registerMainOutputFrom(beamNode, vertex, output)); |
| } |
| |
| /** |
| * @param ctx provides translation context |
| * @param beamNode the beam node to be translated |
| * @param transform transform which can be obtained from {@code beamNode} |
| */ |
| @PrimitiveTransformTranslator(ParDo.SingleOutput.class) |
| private static void parDoSingleOutputTranslator(final PipelineTranslationContext ctx, |
| final TransformHierarchy.Node beamNode, |
| final ParDo.SingleOutput<?, ?> transform) { |
| final Map<Integer, PCollectionView<?>> sideInputMap = getSideInputMap(transform.getSideInputs().values()); |
| final AbstractDoFnTransform doFnTransform = createDoFnTransform(ctx, beamNode, sideInputMap); |
| final IRVertex vertex = new OperatorVertex(doFnTransform); |
| |
| ctx.addVertex(vertex); |
| beamNode.getInputs().values().stream() |
| .filter(input -> !transform.getAdditionalInputs().values().contains(input)) |
| .forEach(input -> ctx.addEdgeTo(vertex, input)); |
| ctx.addSideInputEdges(vertex, sideInputMap); |
| beamNode.getOutputs().values().forEach(output -> ctx.registerMainOutputFrom(beamNode, vertex, output)); |
| } |
| |
| /** |
| * @param ctx provides translation context |
| * @param beamNode the beam node to be translated |
| * @param transform transform which can be obtained from {@code beamNode} |
| */ |
| @PrimitiveTransformTranslator(ParDo.MultiOutput.class) |
| private static void parDoMultiOutputTranslator(final PipelineTranslationContext ctx, |
| final TransformHierarchy.Node beamNode, |
| final ParDo.MultiOutput<?, ?> transform) { |
| final Map<Integer, PCollectionView<?>> sideInputMap = getSideInputMap(transform.getSideInputs().values()); |
| final AbstractDoFnTransform doFnTransform = createDoFnTransform(ctx, beamNode, sideInputMap); |
| final IRVertex vertex = new OperatorVertex(doFnTransform); |
| ctx.addVertex(vertex); |
| beamNode.getInputs().values().stream() |
| .filter(input -> !transform.getAdditionalInputs().values().contains(input)) |
| .forEach(input -> ctx.addEdgeTo(vertex, input)); |
| ctx.addSideInputEdges(vertex, sideInputMap); |
| beamNode.getOutputs().entrySet().stream() |
| .filter(pValueWithTupleTag -> pValueWithTupleTag.getKey().equals(transform.getMainOutputTag())) |
| .forEach(pValueWithTupleTag -> ctx.registerMainOutputFrom(beamNode, vertex, pValueWithTupleTag.getValue())); |
| beamNode.getOutputs().entrySet().stream() |
| .filter(pValueWithTupleTag -> !pValueWithTupleTag.getKey().equals(transform.getMainOutputTag())) |
| .forEach(pValueWithTupleTag -> ctx.registerAdditionalOutputFrom(beamNode, vertex, pValueWithTupleTag.getValue(), |
| pValueWithTupleTag.getKey())); |
| } |
| |
| /** |
| * @param ctx provides translation context |
| * @param beamNode the beam node to be translated |
| * @param transform transform which can be obtained from {@code beamNode} |
| */ |
| @PrimitiveTransformTranslator(GroupByKey.class) |
| private static void groupByKeyTranslator(final PipelineTranslationContext ctx, |
| final TransformHierarchy.Node beamNode, |
| final GroupByKey<?, ?> transform) { |
| final IRVertex vertex = new OperatorVertex(createGBKTransform(ctx, beamNode)); |
| ctx.addVertex(vertex); |
| beamNode.getInputs().values().forEach(input -> ctx.addEdgeTo(vertex, input)); |
| beamNode.getOutputs().values().forEach(output -> ctx.registerMainOutputFrom(beamNode, vertex, output)); |
| } |
| |
| /** |
| * @param ctx provides translation context |
| * @param beamNode the beam node to be translated |
| * @param transform transform which can be obtained from {@code beamNode} |
| */ |
| @PrimitiveTransformTranslator({Window.class, Window.Assign.class}) |
| private static void windowTranslator(final PipelineTranslationContext ctx, |
| final TransformHierarchy.Node beamNode, |
| final PTransform<?, ?> transform) { |
| final WindowFn windowFn; |
| if (transform instanceof Window) { |
| windowFn = ((Window) transform).getWindowFn(); |
| } else if (transform instanceof Window.Assign) { |
| windowFn = ((Window.Assign) transform).getWindowFn(); |
| } else { |
| throw new UnsupportedOperationException(String.format("%s is not supported", transform)); |
| } |
| final IRVertex vertex = new OperatorVertex( |
| new WindowFnTransform(windowFn, DisplayData.from(beamNode.getTransform()))); |
| ctx.addVertex(vertex); |
| beamNode.getInputs().values().forEach(input -> ctx.addEdgeTo(vertex, input)); |
| beamNode.getOutputs().values().forEach(output -> ctx.registerMainOutputFrom(beamNode, vertex, output)); |
| } |
| |
| /** |
| * @param ctx provides translation context |
| * @param beamNode the beam node to be translated |
| * @param transform transform which can be obtained from {@code beamNode} |
| */ |
| @PrimitiveTransformTranslator(View.CreatePCollectionView.class) |
| private static void createPCollectionViewTranslator(final PipelineTranslationContext ctx, |
| final TransformHierarchy.Node beamNode, |
| final View.CreatePCollectionView<?, ?> transform) { |
| final IRVertex vertex = new OperatorVertex(new CreateViewTransform(transform.getView().getViewFn())); |
| ctx.addVertex(vertex); |
| beamNode.getInputs().values().forEach(input -> ctx.addEdgeTo(vertex, input)); |
| ctx.registerMainOutputFrom(beamNode, vertex, transform.getView()); |
| beamNode.getOutputs().values().forEach(output -> ctx.registerMainOutputFrom(beamNode, vertex, output)); |
| } |
| |
| /** |
| * @param ctx provides translation context |
| * @param beamNode the beam node to be translated |
| * @param transform transform which can be obtained from {@code beamNode} |
| */ |
| @PrimitiveTransformTranslator(Flatten.PCollections.class) |
| private static void flattenTranslator(final PipelineTranslationContext ctx, |
| final TransformHierarchy.Node beamNode, |
| final Flatten.PCollections<?> transform) { |
| final IRVertex vertex = new OperatorVertex(new FlattenTransform()); |
| ctx.addVertex(vertex); |
| beamNode.getInputs().values().forEach(input -> ctx.addEdgeTo(vertex, input)); |
| beamNode.getOutputs().values().forEach(output -> ctx.registerMainOutputFrom(beamNode, vertex, output)); |
| } |
| |
| //////////////////////////////////////////////////////////////////////////////////////////////////////// |
| /////////////////////// COMPOSITE TRANSFORMS |
| |
| /** |
| * {@link Combine.PerKey} = {@link GroupByKey} + {@link Combine.GroupedValues} |
| * ({@link Combine.Globally} internally uses {@link Combine.PerKey} which will also be optimized by this translator) |
| * Here, we translate this composite transform as a whole, exploiting its accumulator semantics. |
| * |
| * @param ctx provides translation context |
| * @param beamNode the beam node to be translated |
| * @param transform transform which can be obtained from {@code beamNode} |
| * @return behavior controls whether or not child transforms are visited. |
| */ |
| @CompositeTransformTranslator(Combine.PerKey.class) |
| private static Pipeline.PipelineVisitor.CompositeBehavior combinePerKeyTranslator( |
| final PipelineTranslationContext ctx, |
| final TransformHierarchy.Node beamNode, |
| final PTransform<?, ?> transform) { |
| |
| final Combine.PerKey perKey = (Combine.PerKey) transform; |
| |
| // If there's any side inputs, translate each primitive transforms in this composite transform one by one. |
| if (!perKey.getSideInputs().isEmpty()) { |
| // TODO #264: Partial Combining with Beam SideInputs |
| return Pipeline.PipelineVisitor.CompositeBehavior.ENTER_TRANSFORM; |
| } |
| |
| final CombineFnBase.GlobalCombineFn combineFn = perKey.getFn(); |
| final PCollection<?> mainInput = (PCollection<?>) Iterables.getOnlyElement( |
| TransformInputs.nonAdditionalInputs(beamNode.toAppliedPTransform(ctx.getPipeline()))); |
| final PCollection inputs = (PCollection) Iterables.getOnlyElement( |
| TransformInputs.nonAdditionalInputs(beamNode.toAppliedPTransform(ctx.getPipeline()))); |
| final KvCoder inputCoder = (KvCoder) inputs.getCoder(); |
| final Coder accumulatorCoder; |
| |
| // Check if accumulator coder exists |
| try { |
| accumulatorCoder = |
| combineFn.getAccumulatorCoder(ctx.getPipeline().getCoderRegistry(), inputCoder.getValueCoder()); |
| } catch (CannotProvideCoderException e) { |
| throw new RuntimeException(e); |
| } |
| |
| // If there's no side inputs, |
| // this Combine can be optimized as the following sequence of Nemo IRVertices. |
| // Combine Input -> partialCombine -> finalCombine -> Combine Output |
| final IRVertex partialCombine; |
| final IRVertex finalCombine; |
| |
| // Choose between batch processing and stream processing based on window type and boundedness of data |
| if (isMainInputBounded(beamNode, ctx.getPipeline()) && isGlobalWindow(beamNode, ctx.getPipeline())) { |
| // Batch processing, using CombinePartialTransform and CombineFinalTransform |
| partialCombine = new OperatorVertex(new CombineFnPartialTransform<>(combineFn)); |
| finalCombine = new OperatorVertex(new CombineFnFinalTransform<>(combineFn)); |
| } else { |
| // Stream data processing, using GBKTransform |
| final AppliedPTransform pTransform = beamNode.toAppliedPTransform(ctx.getPipeline()); |
| final CombineFnBase.GlobalCombineFn partialCombineFn = new PartialCombineFn( |
| (Combine.CombineFn) combineFn, accumulatorCoder); |
| final CombineFnBase.GlobalCombineFn finalCombineFn = new FinalCombineFn( |
| (Combine.CombineFn) combineFn, accumulatorCoder); |
| final SystemReduceFn partialSystemReduceFn = |
| SystemReduceFn.combining( |
| inputCoder.getKeyCoder(), |
| AppliedCombineFn.withInputCoder(partialCombineFn, |
| ctx.getPipeline().getCoderRegistry(), inputCoder, |
| null, |
| mainInput.getWindowingStrategy())); |
| final SystemReduceFn finalSystemReduceFn = |
| SystemReduceFn.combining( |
| inputCoder.getKeyCoder(), |
| AppliedCombineFn.withInputCoder(finalCombineFn, |
| ctx.getPipeline().getCoderRegistry(), |
| KvCoder.of(inputCoder.getKeyCoder(), |
| accumulatorCoder), |
| null, mainInput.getWindowingStrategy())); |
| final TupleTag<?> partialMainOutputTag = new TupleTag<>(); |
| final GBKTransform partialCombineStreamTransform = |
| new GBKTransform(inputCoder, |
| Collections.singletonMap(partialMainOutputTag, KvCoder.of(inputCoder.getKeyCoder(), accumulatorCoder)), |
| partialMainOutputTag, |
| mainInput.getWindowingStrategy(), |
| ctx.getPipelineOptions(), |
| partialSystemReduceFn, |
| DoFnSchemaInformation.create(), |
| DisplayData.from(beamNode.getTransform()), |
| true); |
| |
| final GBKTransform finalCombineStreamTransform = |
| new GBKTransform(KvCoder.of(inputCoder.getKeyCoder(), accumulatorCoder), |
| getOutputCoders(pTransform), |
| Iterables.getOnlyElement(beamNode.getOutputs().keySet()), |
| mainInput.getWindowingStrategy(), |
| ctx.getPipelineOptions(), |
| finalSystemReduceFn, |
| DoFnSchemaInformation.create(), |
| DisplayData.from(beamNode.getTransform()), |
| false); |
| |
| partialCombine = new OperatorVertex(partialCombineStreamTransform); |
| finalCombine = new OperatorVertex(finalCombineStreamTransform); |
| } |
| |
| // (Step 1) Partial Combine |
| ctx.addVertex(partialCombine); |
| beamNode.getInputs().values().forEach(input -> ctx.addEdgeTo(partialCombine, input)); |
| |
| // (Step 2) Final Combine |
| ctx.addVertex(finalCombine); |
| beamNode.getOutputs().values().forEach(output -> ctx.registerMainOutputFrom(beamNode, finalCombine, output)); |
| |
| // (Step 3) Adding an edge from partialCombine vertex to finalCombine vertex |
| final IREdge edge = new IREdge(CommunicationPatternProperty.Value.SHUFFLE, partialCombine, finalCombine); |
| final Coder intermediateCoder = KvCoder.of(inputCoder.getKeyCoder(), accumulatorCoder); |
| ctx.addEdge(edge, intermediateCoder, mainInput.getWindowingStrategy().getWindowFn().windowCoder()); |
| |
| // This composite transform has been translated in its entirety. |
| return Pipeline.PipelineVisitor.CompositeBehavior.DO_NOT_ENTER_TRANSFORM; |
| } |
| |
| /** |
| * @param ctx provides translation context |
| * @param beamNode the beam node to be translated |
| * @param transform transform which can be obtained from {@code beamNode} |
| * @return behavior controls whether or not child transforms are visited. |
| */ |
| @CompositeTransformTranslator(LoopCompositeTransform.class) |
| private static Pipeline.PipelineVisitor.CompositeBehavior loopTranslator( |
| final PipelineTranslationContext ctx, |
| final TransformHierarchy.Node beamNode, |
| final LoopCompositeTransform<?, ?> transform) { |
| // Do nothing here, as the context handles the loop vertex stack. |
| // We just keep this method to signal that the loop vertex is acknowledged. |
| return Pipeline.PipelineVisitor.CompositeBehavior.ENTER_TRANSFORM; |
| } |
| |
| //////////////////////////////////////////////////////////////////////////////////////////////////////// |
| /////////////////////// HELPER METHODS |
| |
| /** |
| * @param viewList list of {@link PCollectionView}s. |
| * @return map of side inputs. |
| */ |
| private static Map<Integer, PCollectionView<?>> getSideInputMap(final Collection<PCollectionView<?>> viewList) { |
| final PrimitiveIterator.OfInt iterator = IntStream.range(0, viewList.size()).iterator(); |
| return viewList.stream().collect(Collectors.toMap(i -> iterator.next(), Function.identity())); |
| } |
| |
| /** |
| * @param ctx provides translation context. |
| * @param beamNode the beam node to be translated. |
| * @param sideInputMap side inputs. |
| * @return the created DoFnTransform. |
| */ |
| private static AbstractDoFnTransform createDoFnTransform(final PipelineTranslationContext ctx, |
| final TransformHierarchy.Node beamNode, |
| final Map<Integer, PCollectionView<?>> sideInputMap) { |
| try { |
| final AppliedPTransform pTransform = beamNode.toAppliedPTransform(ctx.getPipeline()); |
| final DoFn doFn = ParDoTranslation.getDoFn(pTransform); |
| final TupleTag mainOutputTag = ParDoTranslation.getMainOutputTag(pTransform); |
| final TupleTagList additionalOutputTags = ParDoTranslation.getAdditionalOutputTags(pTransform); |
| |
| final PCollection<?> mainInput = (PCollection<?>) |
| Iterables.getOnlyElement(TransformInputs.nonAdditionalInputs(pTransform)); |
| |
| final HasDisplayData displayData = (builder) -> builder.add(DisplayData.item("name", beamNode.getFullName())); |
| final DoFnSchemaInformation doFnSchemaInformation = |
| ParDoTranslation.getSchemaInformation(beamNode.toAppliedPTransform(ctx.getPipeline())); |
| |
| if (sideInputMap.isEmpty()) { |
| return new DoFnTransform( |
| doFn, |
| mainInput.getCoder(), |
| getOutputCoders(pTransform), |
| mainOutputTag, |
| additionalOutputTags.getAll(), |
| mainInput.getWindowingStrategy(), |
| ctx.getPipelineOptions(), |
| DisplayData.from(displayData), |
| doFnSchemaInformation, |
| Collections.emptyMap()); |
| } else { |
| return new PushBackDoFnTransform( |
| doFn, |
| mainInput.getCoder(), |
| getOutputCoders(pTransform), |
| mainOutputTag, |
| additionalOutputTags.getAll(), |
| mainInput.getWindowingStrategy(), |
| sideInputMap, |
| ctx.getPipelineOptions(), |
| DisplayData.from(displayData), |
| doFnSchemaInformation, |
| Collections.emptyMap()); |
| } |
| } catch (final IOException e) { |
| throw new RuntimeException(e); |
| } |
| } |
| |
| /** |
| * @param ptransform PTransform to get coders for its outputs |
| * @return the output coders. |
| */ |
| private static Map<TupleTag<?>, Coder<?>> getOutputCoders(final AppliedPTransform<?, ?, ?> ptransform) { |
| return ptransform |
| .getOutputs() |
| .entrySet() |
| .stream() |
| .filter(e -> e.getValue() instanceof PCollection) |
| .collect(Collectors.toMap(Map.Entry::getKey, e -> ((PCollection) e.getValue()).getCoder())); |
| } |
| |
| /** |
| * Returns the correct type of GroupByKey transform by checking whether global windowing strategy is used. |
| * |
| * @param ctx translation context |
| * @param beamNode the beam node to be translated |
| * @return GroupByKey transform |
| */ |
| private static Transform createGBKTransform( |
| final PipelineTranslationContext ctx, |
| final TransformHierarchy.Node beamNode) { |
| final AppliedPTransform<?, ?, ?> pTransform = beamNode.toAppliedPTransform(ctx.getPipeline()); |
| final PCollection<?> mainInput = (PCollection<?>) |
| Iterables.getOnlyElement(TransformInputs.nonAdditionalInputs(pTransform)); |
| final TupleTag mainOutputTag = Iterables.getOnlyElement(beamNode.getOutputs().keySet()); |
| |
| if (isGlobalWindow(beamNode, ctx.getPipeline())) { |
| // GroupByKey Transform when using a global windowing strategy. |
| return new GroupByKeyTransform(); |
| } else { |
| // GroupByKey Transform when using a non-global windowing strategy. |
| return new GBKTransform<>( |
| (KvCoder) mainInput.getCoder(), |
| getOutputCoders(pTransform), |
| mainOutputTag, |
| mainInput.getWindowingStrategy(), |
| ctx.getPipelineOptions(), |
| SystemReduceFn.buffering(mainInput.getCoder()), |
| DoFnSchemaInformation.create(), |
| DisplayData.from(beamNode.getTransform()), |
| false); |
| } |
| } |
| |
| |
| /** |
| * @param beamNode the beam node to be translated. |
| * @param pipeline pipeline. |
| * @return true if the main input has global window. |
| */ |
| private static boolean isGlobalWindow(final TransformHierarchy.Node beamNode, final Pipeline pipeline) { |
| final AppliedPTransform pTransform = beamNode.toAppliedPTransform(pipeline); |
| final PCollection<?> mainInput = (PCollection<?>) |
| Iterables.getOnlyElement(TransformInputs.nonAdditionalInputs(pTransform)); |
| return mainInput.getWindowingStrategy().getWindowFn() instanceof GlobalWindows; |
| } |
| |
| /** |
| * @param beamNode the beam node to be translated. |
| * @param pipeline pipeline. |
| * @return true if the main input bounded. |
| */ |
| private static boolean isMainInputBounded(final TransformHierarchy.Node beamNode, final Pipeline pipeline) { |
| final AppliedPTransform pTransform = beamNode.toAppliedPTransform(pipeline); |
| final PCollection<?> mainInput = (PCollection<?>) |
| Iterables.getOnlyElement(TransformInputs.nonAdditionalInputs(pTransform)); |
| return mainInput.isBounded() == PCollection.IsBounded.BOUNDED; |
| } |
| } |
