Examples with ReplacementOutput org.apache.beam.sdk.runners.PTransformOverrideFactory.ReplacementOutput used on opensource projects

Example 6 with ReplacementOutput

use of org.apache.beam.sdk.runners.PTransformOverrideFactory.ReplacementOutput in project beam by apache.

the class ReplacementOutputsTest method taggedSucceeds.

@Test
public void taggedSucceeds() {
    PCollectionTuple original = PCollectionTuple.of(intsTag, ints).and(strsTag, strs).and(moreIntsTag, moreInts);
    Map<PCollection<?>, ReplacementOutput> replacements = ReplacementOutputs.tagged(PValues.expandOutput((POutput) original), PCollectionTuple.of(strsTag, replacementStrs).and(moreIntsTag, moreReplacementInts).and(intsTag, replacementInts));
    assertThat(replacements.keySet(), Matchers.containsInAnyOrder(replacementStrs, replacementInts, moreReplacementInts));
    ReplacementOutput intsReplacement = replacements.get(replacementInts);
    ReplacementOutput strsReplacement = replacements.get(replacementStrs);
    ReplacementOutput moreIntsReplacement = replacements.get(moreReplacementInts);
    assertThat(intsReplacement, equalTo(ReplacementOutput.of(TaggedPValue.of(intsTag, ints), TaggedPValue.of(intsTag, replacementInts))));
    assertThat(strsReplacement, equalTo(ReplacementOutput.of(TaggedPValue.of(strsTag, strs), TaggedPValue.of(strsTag, replacementStrs))));
    assertThat(moreIntsReplacement, equalTo(ReplacementOutput.of(TaggedPValue.of(moreIntsTag, moreInts), TaggedPValue.of(moreIntsTag, moreReplacementInts))));
}