Examples with UnionCoder org.apache.beam.sdk.transforms.join.UnionCoder used on opensource projects

Example 1 with UnionCoder

use of org.apache.beam.sdk.transforms.join.UnionCoder in project beam by apache.

the class FlinkStreamingTransformTranslators method transformSideInputs.

private static Tuple2<Map<Integer, PCollectionView<?>>, DataStream<RawUnionValue>> transformSideInputs(Collection<PCollectionView<?>> sideInputs, FlinkStreamingTranslationContext context) {
    // collect all side inputs
    Map<TupleTag<?>, Integer> tagToIntMapping = new HashMap<>();
    Map<Integer, PCollectionView<?>> intToViewMapping = new HashMap<>();
    int count = 0;
    for (PCollectionView<?> sideInput : sideInputs) {
        TupleTag<?> tag = sideInput.getTagInternal();
        intToViewMapping.put(count, sideInput);
        tagToIntMapping.put(tag, count);
        count++;
        Coder<Iterable<WindowedValue<?>>> coder = sideInput.getCoderInternal();
    }
    List<Coder<?>> inputCoders = new ArrayList<>();
    for (PCollectionView<?> sideInput : sideInputs) {
        DataStream<Object> sideInputStream = context.getInputDataStream(sideInput);
        TypeInformation<Object> tpe = sideInputStream.getType();
        if (!(tpe instanceof CoderTypeInformation)) {
            throw new IllegalStateException("Input Stream TypeInformation is no CoderTypeInformation.");
        }
        Coder<?> coder = ((CoderTypeInformation) tpe).getCoder();
        inputCoders.add(coder);
    }
    UnionCoder unionCoder = UnionCoder.of(inputCoders);
    CoderTypeInformation<RawUnionValue> unionTypeInformation = new CoderTypeInformation<>(unionCoder);
    // transform each side input to RawUnionValue and union them
    DataStream<RawUnionValue> sideInputUnion = null;
    for (PCollectionView<?> sideInput : sideInputs) {
        TupleTag<?> tag = sideInput.getTagInternal();
        final int intTag = tagToIntMapping.get(tag);
        DataStream<Object> sideInputStream = context.getInputDataStream(sideInput);
        DataStream<RawUnionValue> unionValueStream = sideInputStream.map(new ToRawUnion<>(intTag)).returns(unionTypeInformation);
        if (sideInputUnion == null) {
            sideInputUnion = unionValueStream;
        } else {
            sideInputUnion = sideInputUnion.union(unionValueStream);
        }
    }
    if (sideInputUnion == null) {
        throw new IllegalStateException("No unioned side inputs, this indicates a bug.");
    }
    return new Tuple2<>(intToViewMapping, sideInputUnion);
}