Examples with DoFnSignature org.apache.beam.sdk.transforms.reflect.DoFnSignature used on opensource projects

Example 26 with DoFnSignature

use of org.apache.beam.sdk.transforms.reflect.DoFnSignature in project twister2 by DSC-SPIDAL.

the class ParDoMultiOutputTranslatorBatch method translateNode.

@Override
public void translateNode(ParDo.MultiOutput<IT, OT> transform, Twister2BatchTranslationContext context) {
    DoFn<IT, OT> doFn;
    doFn = transform.getFn();
    BatchTSetImpl<WindowedValue<IT>> inputTTSet = context.getInputDataSet(context.getInput(transform));
    WindowingStrategy<?, ?> windowingStrategy = context.getInput(transform).getWindowingStrategy();
    Coder<IT> inputCoder = (Coder<IT>) context.getInput(transform).getCoder();
    Map<TupleTag<?>, PValue> outputs = context.getOutputs();
    Map<TupleTag<?>, Coder<?>> outputCoders = context.getOutputCoders();
    DoFnSignature signature = DoFnSignatures.getSignature(transform.getFn().getClass());
    DoFnSchemaInformation doFnSchemaInformation;
    doFnSchemaInformation = ParDoTranslation.getSchemaInformation(context.getCurrentTransform());
    TupleTag<OT> mainOutput = transform.getMainOutputTag();
    List<TupleTag<?>> additionalOutputTags = new ArrayList<>(outputs.size() - 1);
    Collection<PCollectionView<?>> sideInputs = transform.getSideInputs();
    // 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) {
        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();
    // put the main output at index 0, FlinkMultiOutputDoFnFunction  expects this
    outputMap.put(mainOutputTag, 0);
    int count = 1;
    for (TupleTag<?> tag : outputs.keySet()) {
        if (!outputMap.containsKey(tag)) {
            outputMap.put(tag, count++);
        }
    }
    ComputeTSet<RawUnionValue> outputTSet = inputTTSet.direct().<RawUnionValue>compute(new DoFnFunction<OT, IT>(context, doFn, inputCoder, outputCoders, additionalOutputTags, windowingStrategy, sideInputStrategies, mainOutput, doFnSchemaInformation, outputMap));
    for (Map.Entry<TupleTag<?>, PValue> output : outputs.entrySet()) {
        ComputeTSet<WindowedValue<OT>> tempTSet = outputTSet.direct().compute(new OutputTagFilter<>(outputMap.get(output.getKey())));
        context.setOutputDataSet((PCollection) output.getValue(), tempTSet);
    }
}