How to use

PTransformOverrideFactory$ReplacementOutput

in

org.apache.beam.sdk.runners

ReplacementOutput mapping = originalToReplacement.get(output.getValue());
if (mapping != null) {
 if (this.equals(producers.get(mapping.getReplacement().getValue()))) {
  producerInput.remove(mapping.getReplacement().getValue());
  producers.remove(mapping.getReplacement().getValue());
  producers.put(mapping.getOriginal().getValue(), this);
   mapping.getReplacement(),
   mapping.getOriginal());
 newOutputsBuilder.put(output.getKey(), mapping.getOriginal().getValue());
} else {
 newOutputsBuilder.put(output);

resultBuilder.put(
  replacementValue.getValue(),
  ReplacementOutput.of(
    mapped, TaggedPValue.of(replacementValue.getKey(), replacementValue.getValue())));
missingTags.remove(replacementValue.getKey());

Collections.singletonMap(
  replacementOutput.get(longs),
  ReplacementOutput.of(
    TaggedPValue.ofExpandedValue(output),
    TaggedPValue.of(replacementLongs.getKey(), replacementLongs.getValue()))));

Collections.singletonMap(
  replacementOutput.get(longs),
  ReplacementOutput.of(
    TaggedPValue.ofExpandedValue(output),
    TaggedPValue.of(replacementLongs.getKey(), replacementLongs.getValue()))));

  Collections.singletonMap(
    replacementOutput,
    ReplacementOutput.of(TaggedPValue.ofExpandedValue(originalOutput), taggedReplacement));
hierarchy.replaceOutputs(replacementOutputs);

  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))));

private void testStreamingWriteOverride(PipelineOptions options, int expectedNumShards) {
 TestPipeline p = TestPipeline.fromOptions(options);
 StreamingShardedWriteFactory<Object, Void, Object> factory =
   new StreamingShardedWriteFactory<>(p.getOptions());
 WriteFiles<Object, Void, Object> original = WriteFiles.to(new TestSink(tmpFolder.toString()));
 PCollection<Object> objs = (PCollection) p.apply(Create.empty(VoidCoder.of()));
 AppliedPTransform<PCollection<Object>, WriteFilesResult<Void>, WriteFiles<Object, Void, Object>>
   originalApplication =
     AppliedPTransform.of("writefiles", objs.expand(), Collections.emptyMap(), original, p);
 WriteFiles<Object, Void, Object> replacement =
   (WriteFiles<Object, Void, Object>)
     factory.getReplacementTransform(originalApplication).getTransform();
 assertThat(replacement, not(equalTo((Object) original)));
 assertThat(replacement.getNumShardsProvider().get(), equalTo(expectedNumShards));
 WriteFilesResult<Void> originalResult = objs.apply(original);
 WriteFilesResult<Void> replacementResult = objs.apply(replacement);
 Map<PValue, ReplacementOutput> res =
   factory.mapOutputs(originalResult.expand(), replacementResult);
 assertEquals(1, res.size());
 assertEquals(
   originalResult.getPerDestinationOutputFilenames(),
   res.get(replacementResult.getPerDestinationOutputFilenames()).getOriginal().getValue());
}

@Test
public void testMapOutputs() {
 PCollection<Integer> input = pipeline.apply(Create.of(1, 2, 3));
 PCollection<Integer> output = input.apply("Map", MapElements.via(fn));
 PCollection<Integer> reappliedOutput = input.apply("ReMap", MapElements.via(fn));
 Map<PValue, ReplacementOutput> replacementMap =
   factory.mapOutputs(output.expand(), reappliedOutput);
 assertThat(
   replacementMap,
   Matchers.hasEntry(
     reappliedOutput,
     ReplacementOutput.of(
       TaggedPValue.ofExpandedValue(output),
       TaggedPValue.ofExpandedValue(reappliedOutput))));
}

 @Test
 public void outputMapping() {
  final PCollectionList<String> inputList =
    PCollectionList.of(first).and(second).and(first).and(first);
  PCollection<String> original = inputList.apply(Flatten.pCollections());
  PCollection<String> replacement = inputList.apply(new FlattenWithoutDuplicateInputs<>());

  assertThat(
    factory.mapOutputs(original.expand(), replacement),
    Matchers.hasEntry(
      replacement,
      ReplacementOutput.of(
        TaggedPValue.ofExpandedValue(original),
        TaggedPValue.ofExpandedValue(replacement))));
 }
}

@Test
public void mapOutputsSucceeds() {
 PCollection<Long> original = pipeline.apply("Original", GenerateSequence.from(0));
 PCollection<Long> replacement = pipeline.apply("Replacement", GenerateSequence.from(0));
 Map<PValue, ReplacementOutput> mapping = factory.mapOutputs(original.expand(), replacement);
 assertThat(
   mapping,
   Matchers.hasEntry(
     replacement,
     ReplacementOutput.of(
       TaggedPValue.ofExpandedValue(original),
       TaggedPValue.ofExpandedValue(replacement))));
}

 @Override
 public Map<PValue, ReplacementOutput> mapOutputs(
   Map<TupleTag<?>, PValue> outputs, PCollection<T> newOutput) {
  Map.Entry<TupleTag<?>, PValue> original = Iterables.getOnlyElement(outputs.entrySet());
  Map.Entry<TupleTag<?>, PValue> replacement =
    Iterables.getOnlyElement(newOutput.expand().entrySet());
  return Collections.singletonMap(
    newOutput,
    ReplacementOutput.of(
      TaggedPValue.of(original.getKey(), original.getValue()),
      TaggedPValue.of(replacement.getKey(), replacement.getValue())));
 }
}

 @Override
 public Map<PValue, ReplacementOutput> mapOutputs(
   Map<TupleTag<?>, PValue> outputs, PCollection<Long> newOutput) {
  Map.Entry<TupleTag<?>, PValue> original = Iterables.getOnlyElement(outputs.entrySet());
  Map.Entry<TupleTag<?>, PValue> replacement =
    Iterables.getOnlyElement(newOutput.expand().entrySet());
  return Collections.singletonMap(
    newOutput,
    ReplacementOutput.of(
      TaggedPValue.of(original.getKey(), original.getValue()),
      TaggedPValue.of(replacement.getKey(), replacement.getValue())));
 }
}

public static Map<PValue, ReplacementOutput> singleton(
  Map<TupleTag<?>, PValue> original, PValue replacement) {
 Entry<TupleTag<?>, PValue> originalElement = Iterables.getOnlyElement(original.entrySet());
 TupleTag<?> replacementTag = Iterables.getOnlyElement(replacement.expand().entrySet()).getKey();
 return Collections.singletonMap(
   replacement,
   ReplacementOutput.of(
     TaggedPValue.of(originalElement.getKey(), originalElement.getValue()),
     TaggedPValue.of(replacementTag, replacement)));
}

@Test
public void singletonSucceeds() {
 Map<PValue, ReplacementOutput> replacements =
   ReplacementOutputs.singleton(ints.expand(), replacementInts);
 assertThat(replacements, Matchers.hasKey(replacementInts));
 ReplacementOutput replacement = replacements.get(replacementInts);
 Map.Entry<TupleTag<?>, PValue> taggedInts = Iterables.getOnlyElement(ints.expand().entrySet());
 assertThat(replacement.getOriginal().getTag(), equalTo(taggedInts.getKey()));
 assertThat(replacement.getOriginal().getValue(), equalTo(taggedInts.getValue()));
 assertThat(replacement.getReplacement().getValue(), equalTo(replacementInts));
}