PrimitiveParDoSingleFactory extends SingleInputOutputOverrideFactory<
        PCollection<? extends InputT>, PCollection<OutputT>, ParDo.SingleOutput<InputT, OutputT>> { @Override public PTransformReplacement<PCollection<? extends InputT>, PCollection<OutputT>> getReplacementTransform( AppliedPTransform< PCollection<? extends InputT>, PCollection<OutputT>, SingleOutput<InputT, OutputT>> transform) { return PTransformReplacement.of( PTransformReplacements.getSingletonMainInput(transform), new ParDoSingle<>( transform.getTransform(), Iterables.getOnlyElement(transform.getOutputs().keySet()), PTransformReplacements.getSingletonMainOutput(transform).getCoder())); }  @Override PTransformReplacement<PCollection<? extends InputT>, PCollection<OutputT>> getReplacementTransform(
          AppliedPTransform<
                  PCollection<? extends InputT>,
                  PCollection<OutputT>,
                  SingleOutput<InputT, OutputT>>
              transform);  }
@Test public void getReplacementTransformPopulateDisplayData() { ParDo.SingleOutput<Integer, Long> originalTransform = ParDo.of(new ToLongFn()); DisplayData originalDisplayData = DisplayData.from(originalTransform); PCollection<? extends Integer> input = pipeline.apply(Create.of(1, 2, 3)); AppliedPTransform< PCollection<? extends Integer>, PCollection<Long>, ParDo.SingleOutput<Integer, Long>> application = AppliedPTransform.of( "original", input.expand(), input.apply(originalTransform).expand(), originalTransform, pipeline); PTransformReplacement<PCollection<? extends Integer>, PCollection<Long>> replacement = factory.getReplacementTransform(application); DisplayData replacementDisplayData = DisplayData.from(replacement.getTransform()); assertThat(replacementDisplayData, equalTo(originalDisplayData)); DisplayData primitiveDisplayData = Iterables.getOnlyElement( DisplayDataEvaluator.create() .displayDataForPrimitiveTransforms(replacement.getTransform(), VarIntCoder.of())); assertThat(primitiveDisplayData, equalTo(replacementDisplayData)); }
@Test public void getReplacementTransformGetSideInputs() { PCollectionView<Long> sideLong = pipeline .apply("LongSideInputVals", Create.of(-1L, -2L, -4L)) .apply("SideLongView", Sum.longsGlobally().asSingletonView()); PCollectionView<List<String>> sideStrings = pipeline .apply("StringSideInputVals", Create.of("foo", "bar", "baz")) .apply("SideStringsView", View.asList()); ParDo.SingleOutput<Integer, Long> originalTransform = ParDo.of(new ToLongFn()).withSideInputs(sideLong, sideStrings); PCollection<? extends Integer> input = pipeline.apply(Create.of(1, 2, 3)); AppliedPTransform< PCollection<? extends Integer>, PCollection<Long>, ParDo.SingleOutput<Integer, Long>> application = AppliedPTransform.of( "original", input.expand(), input.apply(originalTransform).expand(), originalTransform, pipeline); PTransformReplacement<PCollection<? extends Integer>, PCollection<Long>> replacementTransform = factory.getReplacementTransform(application); ParDoSingle<Integer, Long> parDoSingle = (ParDoSingle<Integer, Long>) replacementTransform.getTransform(); assertThat(parDoSingle.getSideInputs().values(), containsInAnyOrder(sideStrings, sideLong)); }
@Test public void getReplacementTransformGetFn() { DoFn<Integer, Long> originalFn = new ToLongFn(); ParDo.SingleOutput<Integer, Long> originalTransform = ParDo.of(originalFn); PCollection<? extends Integer> input = pipeline.apply(Create.of(1, 2, 3)); AppliedPTransform< PCollection<? extends Integer>, PCollection<Long>, ParDo.SingleOutput<Integer, Long>> application = AppliedPTransform.of( "original", input.expand(), input.apply(originalTransform).expand(), originalTransform, pipeline); PTransformReplacement<PCollection<? extends Integer>, PCollection<Long>> replacementTransform = factory.getReplacementTransform(application); ParDoSingle<Integer, Long> parDoSingle = (ParDoSingle<Integer, Long>) replacementTransform.getTransform(); assertThat(parDoSingle.getFn(), equalTo(originalTransform.getFn())); assertThat(parDoSingle.getFn(), equalTo(originalFn)); }