@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 CompositeBehavior enterCompositeTransform(Node node) { if (node.getTransform() instanceof StatefulSingleOutputParDo) { batchStatefulDoFn = ((StatefulSingleOutputParDo) node.getTransform()).getOriginalParDo().getFn(); return CompositeBehavior.DO_NOT_ENTER_TRANSFORM; } else if (node.getTransform() instanceof StatefulMultiOutputParDo) { batchStatefulDoFn = ((StatefulMultiOutputParDo) node.getTransform()).getOriginalParDo().getFn(); return CompositeBehavior.DO_NOT_ENTER_TRANSFORM; } else { return CompositeBehavior.ENTER_TRANSFORM; } } }