@Override public PTransformReplacement<PCollection<KV<K, V>>, PCollection<KV<K, V>>> getReplacementTransform( AppliedPTransform<PCollection<KV<K, V>>, PCollection<KV<K, V>>, Reshuffle<K, V>> transform) { return PTransformReplacement.of( PTransformReplacements.getSingletonMainInput(transform), new ReshuffleWithOnlyTrigger<>()); }
@Override public PTransformReplacement<PCollection<InputT>, PCollectionTuple> getReplacementTransform( AppliedPTransform<PCollection<InputT>, PCollectionTuple, ParDo.MultiOutput<InputT, OutputT>> appliedTransform) { return PTransformReplacement.of( PTransformReplacements.getSingletonMainInput(appliedTransform), SplittableParDo.forAppliedParDo(appliedTransform)); }
@Override public PTransformReplacement<PCollection<InputT>, PCollectionTuple> getReplacementTransform( AppliedPTransform<PCollection<InputT>, PCollectionTuple, MultiOutput<InputT, OutputT>> transform) { return PTransformReplacement.of( PTransformReplacements.getSingletonMainInput(transform), forAppliedParDo(transform)); }
@Override public PTransformReplacement<PCollection<InputT>, PCollectionTuple> getReplacementTransform( AppliedPTransform< PCollection<InputT>, PCollectionTuple, MultiOutput<InputT, OutputT>> transform) { return PTransformReplacement.of( PTransformReplacements.getSingletonMainInput(transform), (SplittableParDo<InputT, OutputT, ?>) SplittableParDo.forAppliedParDo(transform)); }
/** * Gets the singleton input of an {@link AppliedPTransform}, ignoring any additional inputs * returned by {@link PTransform#getAdditionalInputs()}. */ public static <T> PCollection<T> getSingletonMainInput( AppliedPTransform<? extends PCollection<? extends T>, ?, ?> application) { return getSingletonMainInput( application.getInputs(), application.getTransform().getAdditionalInputs().keySet()); }
@Override public PTransformReplacement<PCollection<KV<K, InputT>>, PCollection<OutputT>> getReplacementTransform( AppliedPTransform< PCollection<KV<K, InputT>>, PCollection<OutputT>, SingleOutput<KV<K, InputT>, OutputT>> transform) { return PTransformReplacement.of( PTransformReplacements.getSingletonMainInput(transform), new StatefulSingleOutputParDo<>(transform.getTransform(), isFnApi)); }
@Override public PTransformReplacement<PCollection<KV<K, InputT>>, PCollectionTuple> getReplacementTransform( AppliedPTransform< PCollection<KV<K, InputT>>, PCollectionTuple, MultiOutput<KV<K, InputT>, OutputT>> transform) { return PTransformReplacement.of( PTransformReplacements.getSingletonMainInput(transform), new StatefulMultiOutputParDo<>(transform.getTransform(), isFnApi)); }
@Override public PTransformReplacement<PCollection<PubsubMessage>, PDone> getReplacementTransform( AppliedPTransform<PCollection<PubsubMessage>, PDone, PubsubUnboundedSink> transform) { return PTransformReplacement.of( PTransformReplacements.getSingletonMainInput(transform), new StreamingPubsubIOWrite(runner, transform.getTransform())); }
@Override public PTransformReplacement< PCollection<KV<byte[], KV<InputT, RestrictionT>>>, PCollectionTuple> getReplacementTransform( AppliedPTransform< PCollection<KV<byte[], KV<InputT, RestrictionT>>>, PCollectionTuple, ProcessKeyedElements<InputT, OutputT, RestrictionT>> transform) { return PTransformReplacement.of( PTransformReplacements.getSingletonMainInput(transform), new SplittableProcessViaKeyedWorkItems<>(transform.getTransform())); }
@Override public PTransformReplacement<PCollection<T>, PCollection<T>> getReplacementTransform( AppliedPTransform<PCollection<T>, PCollection<T>, CreatePCollectionView<T, T>> transform) { return PTransformReplacement.of( PTransformReplacements.getSingletonMainInput(transform), new StreamingWrapSingletonInList<>(transform.getTransform())); } }
@Override public PTransformReplacement<PCollection<T>, PCollection<T>> getReplacementTransform( AppliedPTransform<PCollection<T>, PCollection<T>, CreatePCollectionView<T, Iterable<T>>> transform) { return PTransformReplacement.of( PTransformReplacements.getSingletonMainInput(transform), new StreamingViewAsIterable<>(transform.getTransform().getView())); } }
@Override public PTransformReplacement<PCollection<ElemT>, PCollection<ElemT>> getReplacementTransform( AppliedPTransform< PCollection<ElemT>, PCollection<ElemT>, CreatePCollectionView<ElemT, ViewT>> transform) { StreamingCreatePCollectionView<ElemT, ViewT> streamingView = new StreamingCreatePCollectionView<>(transform.getTransform().getView()); return PTransformReplacement.of( PTransformReplacements.getSingletonMainInput(transform), streamingView); }
@Override public PTransformReplacement<PCollection<InputT>, PCollection<OutputT>> getReplacementTransform( AppliedPTransform<PCollection<InputT>, PCollection<OutputT>, TransformT> transform) { PTransform<PCollection<InputT>, PCollection<OutputT>> rep = InstanceBuilder.ofType(replacement) .withArg(DataflowRunner.class, runner) .withArg( (Class<TransformT>) transform.getTransform().getClass(), transform.getTransform()) .build(); return PTransformReplacement.of(PTransformReplacements.getSingletonMainInput(transform), rep); } }
@Override public PTransformReplacement<PCollection<? extends Integer>, PCollection<Integer>> getReplacementTransform( AppliedPTransform< PCollection<? extends Integer>, PCollection<Integer>, MapElements<Integer, Integer>> transform) { return PTransformReplacement.of( PTransformReplacements.getSingletonMainInput(transform), transform.getTransform()); } };
@Override public PTransformReplacement<PCollection<KV<K, Iterable<InputT>>>, PCollection<KV<K, OutputT>>> getReplacementTransform( AppliedPTransform< PCollection<KV<K, Iterable<InputT>>>, PCollection<KV<K, OutputT>>, GroupedValues<K, InputT, OutputT>> transform) { return PTransformReplacement.of( PTransformReplacements.getSingletonMainInput(transform), new CombineGroupedValues<>( transform.getTransform(), PTransformReplacements.getSingletonMainOutput(transform).getCoder())); }
@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 public PTransformReplacement<PCollection<InputT>, PValue> getReplacementTransform( final AppliedPTransform< PCollection<InputT>, PValue, PTransform<PCollection<InputT>, PValue>> transform) { PTransform<PCollection<InputT>, PValue> rep = InstanceBuilder.ofType(replacement) .withArg(DataflowRunner.class, runner) .withArg(CreatePCollectionView.class, findCreatePCollectionView(transform)) .build(); return PTransformReplacement.of( PTransformReplacements.getSingletonMainInput(transform), (PTransform) rep); }
@Test public void getMainInputSingleOutputSingleInput() { AppliedPTransform<PCollection<Long>, ?, ?> application = AppliedPTransform.of( "application", Collections.singletonMap(new TupleTag<Long>(), mainInput), Collections.singletonMap(new TupleTag<Long>(), output), ParDo.of(new TestDoFn()), pipeline); PCollection<Long> input = PTransformReplacements.getSingletonMainInput(application); assertThat(input, equalTo(mainInput)); }
@Override public PTransformReplacement< PCollection<KV<byte[], KV<InputT, RestrictionT>>>, PCollectionTuple> getReplacementTransform( AppliedPTransform< PCollection<KV<byte[], KV<InputT, RestrictionT>>>, PCollectionTuple, ProcessKeyedElements<InputT, OutputT, RestrictionT>> transform) { checkArgument( DoFnSignatures.signatureForDoFn(transform.getTransform().getFn()).isBoundedPerElement() == IsBounded.BOUNDED, "Expecting a bounded-per-element splittable DoFn"); return PTransformReplacement.of( PTransformReplacements.getSingletonMainInput(transform), new SplittableProcessNaive<>(transform.getTransform())); }
@Test public void getMainInputSingleOutputSideInputs() { AppliedPTransform<PCollection<Long>, ?, ?> application = AppliedPTransform.of( "application", ImmutableMap.<TupleTag<?>, PValue>builder() .put(new TupleTag<Long>(), mainInput) .put(sideInput.getTagInternal(), sideInput.getPCollection()) .build(), Collections.singletonMap(new TupleTag<Long>(), output), ParDo.of(new TestDoFn()).withSideInputs(sideInput), pipeline); PCollection<Long> input = PTransformReplacements.getSingletonMainInput(application); assertThat(input, equalTo(mainInput)); }