Search in sources :

Example 1 with TupleTagList

use of org.apache.beam.sdk.values.TupleTagList in project beam by apache.

the class Partition method expand.

// ///////////////////////////////////////////////////////////////////////////
@Override
public PCollectionList<T> expand(PCollection<T> in) {
    final TupleTagList outputTags = partitionDoFn.getOutputTags();
    PCollectionTuple outputs = in.apply(ParDo.of(partitionDoFn).withOutputTags(new TupleTag<Void>() {
    }, outputTags).withSideInputs(partitionDoFn.getSideInputs()));
    PCollectionList<T> pcs = PCollectionList.empty(in.getPipeline());
    Coder<T> coder = in.getCoder();
    for (TupleTag<?> outputTag : outputTags.getAll()) {
        // All the tuple tags are actually TupleTag<T>
        // And all the collections are actually PCollection<T>
        @SuppressWarnings("unchecked") TupleTag<T> typedOutputTag = (TupleTag<T>) outputTag;
        pcs = pcs.and(outputs.get(typedOutputTag).setCoder(coder));
    }
    return pcs;
}
Also used : TupleTagList(org.apache.beam.sdk.values.TupleTagList) PCollectionTuple(org.apache.beam.sdk.values.PCollectionTuple) TupleTag(org.apache.beam.sdk.values.TupleTag)

Aggregations

PCollectionTuple (org.apache.beam.sdk.values.PCollectionTuple)1 TupleTag (org.apache.beam.sdk.values.TupleTag)1 TupleTagList (org.apache.beam.sdk.values.TupleTagList)1