ReplacementOutputs { public static Map<PValue, ReplacementOutput> tagged( Map<TupleTag<?>, PValue> original, POutput replacement) { Map<TupleTag<?>, TaggedPValue> originalTags = new HashMap<>(); for (Map.Entry<TupleTag<?>, PValue> originalValue : original.entrySet()) { originalTags.put( originalValue.getKey(), TaggedPValue.of(originalValue.getKey(), originalValue.getValue())); } ImmutableMap.Builder<PValue, ReplacementOutput> resultBuilder = ImmutableMap.builder(); Set<TupleTag<?>> missingTags = new HashSet<>(originalTags.keySet()); for (Map.Entry<TupleTag<?>, PValue> replacementValue : replacement.expand().entrySet()) { TaggedPValue mapped = originalTags.get(replacementValue.getKey()); checkArgument( mapped != null, "Missing original output for Tag %s and Value %s Between original %s and replacement %s", replacementValue.getKey(), replacementValue.getValue(), original, replacement.expand()); resultBuilder.put( replacementValue.getValue(), ReplacementOutput.of( mapped, TaggedPValue.of(replacementValue.getKey(), replacementValue.getValue()))); missingTags.remove(replacementValue.getKey()); } ImmutableMap<PValue, ReplacementOutput> result = resultBuilder.build(); checkArgument( missingTags.isEmpty(), "Missing replacement for tags %s. Encountered tags: %s", missingTags, result.keySet()); return result; } private  ReplacementOutputs(); static Map<PValue, ReplacementOutput> singleton(
      Map<TupleTag<?>, PValue> original, PValue replacement); static Map<PValue, ReplacementOutput> tagged(
      Map<TupleTag<?>, PValue> original, POutput replacement);  }
@Test public void taggedSucceeds() { PCollectionTuple original = PCollectionTuple.of(intsTag, ints).and(strsTag, strs).and(moreIntsTag, moreInts); Map<PValue, ReplacementOutput> replacements = ReplacementOutputs.tagged( original.expand(), 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)))); }
@Test public void taggedMissingReplacementThrows() { PCollectionTuple original = PCollectionTuple.of(intsTag, ints).and(strsTag, strs).and(moreIntsTag, moreInts); thrown.expect(IllegalArgumentException.class); thrown.expectMessage("Missing replacement"); thrown.expectMessage(intsTag.toString()); thrown.expectMessage(ints.toString()); ReplacementOutputs.tagged( original.expand(), PCollectionTuple.of(strsTag, replacementStrs).and(moreIntsTag, moreReplacementInts)); }
@Test public void taggedExtraReplacementThrows() { PCollectionTuple original = PCollectionTuple.of(intsTag, ints).and(strsTag, strs); thrown.expect(IllegalArgumentException.class); thrown.expectMessage("Missing original output"); thrown.expectMessage(moreIntsTag.toString()); thrown.expectMessage(moreReplacementInts.toString()); ReplacementOutputs.tagged( original.expand(), PCollectionTuple.of(strsTag, replacementStrs) .and(moreIntsTag, moreReplacementInts) .and(intsTag, replacementInts)); }