Examples with DoFnFunction org.apache.beam.runners.twister2.translators.functions.DoFnFunction used on opensource projects

Example 1 with DoFnFunction

use of org.apache.beam.runners.twister2.translators.functions.DoFnFunction in project beam by apache.

the class BeamBatchWorker method addInputs.

/**
 * Adds all the side inputs into the sink test so it is available from the DoFn's.
 */
private void addInputs(BaseTSet sinkTSet, Map<String, CachedTSet> sideInputTSets) {
    if (sideInputTSets.isEmpty()) {
        return;
    }
    TBaseGraph graph = sinkTSet.getTBaseGraph();
    TBase currNode = null;
    Deque<TBase> deque = new ArrayDeque<>();
    deque.add(sinkTSet);
    while (!deque.isEmpty()) {
        currNode = deque.remove();
        deque.addAll(graph.getPredecessors(currNode));
        if (currNode instanceof ComputeTSet) {
            if (((ComputeTSet) currNode).getComputeFunc() instanceof DoFnFunction) {
                Set<String> sideInputKeys = ((DoFnFunction) ((ComputeTSet) currNode).getComputeFunc()).getSideInputKeys();
                for (String sideInputKey : sideInputKeys) {
                    if (!sideInputTSets.containsKey(sideInputKey)) {
                        throw new IllegalStateException("Side input not found for key " + sideInputKey);
                    }
                    ((ComputeTSet) currNode).addInput(sideInputKey, sideInputTSets.get(sideInputKey));
                }
            }
        }
    }
}

Example 2 with DoFnFunction

use of org.apache.beam.runners.twister2.translators.functions.DoFnFunction in project twister2 by DSC-SPIDAL.

the class Twister2TranslationContext method addInputs.

/**
 * Adds all the side inputs into the sink test so it is available from the DoFn's
 */
private void addInputs(SinkTSet sinkTSet, Map<String, CachedTSet> sideInputTSets) {
    TBaseGraph graph = sinkTSet.getTBaseGraph();
    Set<String> keys = sideInputTSets.keySet();
    TBase currNode = null;
    Deque<TBase> deque = new ArrayDeque<>();
    deque.add(sinkTSet);
    while (!deque.isEmpty()) {
        currNode = deque.remove();
        deque.addAll(graph.getPredecessors(currNode));
        if (currNode instanceof ComputeTSet) {
            if (((ComputeTSet) currNode).getComputeFunc() instanceof DoFnFunction) {
                Set<String> sideInputKeys = ((DoFnFunction) ((ComputeTSet) currNode).getComputeFunc()).getSideInputKeys();
                for (String sideInputKey : sideInputKeys) {
                    if (!sideInputTSets.containsKey(sideInputKey)) {
                        throw new IllegalStateException("Side input not found for key " + sideInputKey);
                    }
                    ((ComputeTSet) currNode).addInput(sideInputKey, sideInputTSets.get(sideInputKey));
                }
            }
        }
    }
}

Example 3 with DoFnFunction

use of org.apache.beam.runners.twister2.translators.functions.DoFnFunction in project beam by apache.

the class ParDoMultiOutputTranslatorBatch method translateNode.

@Override
public void translateNode(ParDo.MultiOutput<InputT, OutputT> transform, Twister2BatchTranslationContext context) {
    DoFn<InputT, OutputT> doFn;
    doFn = transform.getFn();
    if (DoFnSignatures.signatureForDoFn(doFn).processElement().isSplittable()) {
        throw new UnsupportedOperationException(String.format("Not expected to directly translate splittable DoFn, should have been overridden: %s", doFn));
    }
    BatchTSetImpl<WindowedValue<InputT>> inputTTSet = context.getInputDataSet(context.getInput(transform));
    WindowingStrategy<?, ?> windowingStrategy = context.getInput(transform).getWindowingStrategy();
    Coder<InputT> inputCoder = (Coder<InputT>) context.getInput(transform).getCoder();
    Map<String, PCollectionView<?>> sideInputMapping;
    Map<TupleTag<?>, PCollection<?>> outputs = context.getOutputs();
    Map<TupleTag<?>, Coder<?>> outputCoders = context.getOutputCoders();
    // DoFnSignature signature = DoFnSignatures.getSignature(transform.getFn().getClass());
    DoFnSchemaInformation doFnSchemaInformation;
    doFnSchemaInformation = ParDoTranslation.getSchemaInformation(context.getCurrentTransform());
    sideInputMapping = ParDoTranslation.getSideInputMapping(context.getCurrentTransform());
    TupleTag<OutputT> mainOutput = transform.getMainOutputTag();
    List<TupleTag<?>> additionalOutputTags = new ArrayList<>(transform.getAdditionalOutputTags().getAll());
    Map<String, PCollectionView<?>> sideInputs = transform.getSideInputs();
    // TODO : note change from List to map in sideinputs
    // construct a map from side input to WindowingStrategy so that
    // the DoFn runner can map main-input windows to side input windows
    Map<PCollectionView<?>, WindowingStrategy<?, ?>> sideInputStrategies = new HashMap<>();
    for (PCollectionView<?> sideInput : sideInputs.values()) {
        sideInputStrategies.put(sideInput, sideInput.getWindowingStrategyInternal());
    }
    TupleTag<?> mainOutputTag;
    try {
        mainOutputTag = ParDoTranslation.getMainOutputTag(context.getCurrentTransform());
    } catch (IOException e) {
        throw new RuntimeException(e);
    }
    Map<TupleTag<?>, Integer> outputMap = Maps.newHashMap();
    outputMap.put(mainOutputTag, 0);
    int count = 1;
    for (TupleTag<?> tag : outputs.keySet()) {
        if (!outputMap.containsKey(tag)) {
            outputMap.put(tag, count++);
        }
    }
    ComputeTSet<RawUnionValue, Iterator<WindowedValue<InputT>>> outputTSet = inputTTSet.direct().<RawUnionValue>compute(new DoFnFunction<OutputT, InputT>(context, doFn, inputCoder, outputCoders, additionalOutputTags, windowingStrategy, sideInputStrategies, mainOutput, doFnSchemaInformation, outputMap, sideInputMapping));
    for (Map.Entry<TupleTag<?>, PCollection<?>> output : outputs.entrySet()) {
        ComputeTSet<WindowedValue<OutputT>, Iterator<RawUnionValue>> tempTSet = outputTSet.direct().compute(new OutputTagFilter(outputMap.get(output.getKey())));
        context.setOutputDataSet((PCollection) output.getValue(), tempTSet);
    }
}