@Override public PTransformReplacement<PCollection<InputT>, PCollectionTuple> getReplacementTransform( AppliedPTransform<PCollection<InputT>, PCollectionTuple, ParDo.MultiOutput<InputT, OutputT>> appliedTransform) { return PTransformReplacement.of( PTransformReplacements.getSingletonMainInput(appliedTransform), new PTransform<PCollection<InputT>, PCollectionTuple>() { @Override public PCollectionTuple expand(PCollection<InputT> input) { return input .apply("Materialize input", Reshuffle.viaRandomKey()) .apply("ParDo with stable input", appliedTransform.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<InputT>, PCollection<OutputT>> getReplacementTransform( AppliedPTransform< PCollection<InputT>, PCollection<OutputT>, ParDo.SingleOutput<InputT, OutputT>> appliedTransform) { return PTransformReplacement.of( PTransformReplacements.getSingletonMainInput(appliedTransform), new PTransform<PCollection<InputT>, PCollection<OutputT>>() { @Override public PCollection<OutputT> expand(PCollection<InputT> input) { return input .apply("Materialize input", Reshuffle.viaRandomKey()) .apply("ParDo with stable input", appliedTransform.getTransform()); } }); }
@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<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<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); }
@Override public PTransformReplacement<PCollection<InputT>, PValue> getReplacementTransform( AppliedPTransform<PCollection<InputT>, PValue, PTransform<PCollection<InputT>, PValue>> transform) { Combine.GloballyAsSingletonView<?, ?> combineTransform = (Combine.GloballyAsSingletonView) transform.getTransform(); return PTransformReplacement.of( PTransformReplacements.getSingletonMainInput(transform), new BatchViewOverrides.BatchViewAsSingleton( runner, findCreatePCollectionView(transform), (CombineFn) combineTransform.getCombineFn(), combineTransform.getFanout())); } }