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

Example 1 with ByteToElemFunction

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

the class PCollectionViewTranslatorBatch method translateNode.

@Override
public void translateNode(View.CreatePCollectionView<ElemT, ViewT> transform, Twister2BatchTranslationContext context) {
    BatchTSet<WindowedValue<ElemT>> inputDataSet = context.getInputDataSet(context.getInput(transform));
    @SuppressWarnings("unchecked") AppliedPTransform<PCollection<ElemT>, PCollection<ElemT>, PTransform<PCollection<ElemT>, PCollection<ElemT>>> application = (AppliedPTransform<PCollection<ElemT>, PCollection<ElemT>, PTransform<PCollection<ElemT>, PCollection<ElemT>>>) context.getCurrentTransform();
    org.apache.beam.sdk.values.PCollectionView<ViewT> input;
    PCollection<ElemT> inputPCol = context.getInput(transform);
    final Coder coder = inputPCol.getCoder();
    WindowingStrategy windowingStrategy = inputPCol.getWindowingStrategy();
    WindowFn windowFn = windowingStrategy.getWindowFn();
    try {
        input = CreatePCollectionViewTranslation.getView(application);
    } catch (IOException e) {
        throw new RuntimeException(e);
    }
    switch(input.getViewFn().getMaterialization().getUrn()) {
        case Materializations.MULTIMAP_MATERIALIZATION_URN:
            KvCoder kvCoder = (KvCoder<?, ?>) coder;
            final Coder keyCoder = kvCoder.getKeyCoder();
            final WindowedValue.WindowedValueCoder kvwvCoder = WindowedValue.FullWindowedValueCoder.of(kvCoder.getValueCoder(), windowFn.windowCoder());
            BatchTSet<WindowedValue<ElemT>> multimapMaterialization = inputDataSet.direct().map(new MapToTupleFunction<>(keyCoder, kvwvCoder)).allGather().map(new ByteToWindowFunctionPrimitive(keyCoder, kvwvCoder));
            context.setSideInputDataSet(input.getTagInternal().getId(), multimapMaterialization);
            break;
        case Materializations.ITERABLE_MATERIALIZATION_URN:
            final WindowedValue.WindowedValueCoder wvCoder = WindowedValue.FullWindowedValueCoder.of(coder, windowFn.windowCoder());
            BatchTSet<WindowedValue<ElemT>> iterableMaterialization = inputDataSet.direct().map(new ElemToBytesFunction<>(wvCoder)).allGather().map(new ByteToElemFunction(wvCoder));
            try {
                input = CreatePCollectionViewTranslation.getView(application);
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
            context.setSideInputDataSet(input.getTagInternal().getId(), iterableMaterialization);
            break;
        default:
            throw new UnsupportedOperationException("Unknown side input materialization " + input.getViewFn().getMaterialization().getUrn());
    }
}