stepContext, ptransformId, transform.getFn(), transform.getInputWindowingStrategy(), transform.getSideInputs(),
throws IOException { ProcessKeyedElements<?, ?, ?> pke = transform.getTransform(); final DoFn<?, ?> fn = pke.getFn(); final DoFnSignature signature = DoFnSignatures.getSignature(fn.getClass()); final String restrictionCoderId = components.registerCoder(pke.getRestrictionCoder());
@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())); }
public DoFn<InputT, OutputT> getFn() { return original.getFn(); }
@Override public PCollectionTuple expand(PCollection<KV<byte[], KV<InputT, RestrictionT>>> input) { return input .apply("Drop key", Values.create()) .apply("Reshuffle", Reshuffle.of()) .apply( "NaiveProcess", ParDo.of( new NaiveProcessFn<InputT, OutputT, RestrictionT, TrackerT>(original.getFn())) .withSideInputs(original.getSideInputs()) .withOutputTags(original.getMainOutputTag(), original.getAdditionalOutputTags())); } }
@Override public boolean matches(AppliedPTransform<?, ?, ?> application) { PTransform<?, ?> transform = application.getTransform(); if (transform instanceof SplittableParDo.ProcessKeyedElements) { DoFn<?, ?> fn = ((SplittableParDo.ProcessKeyedElements) transform).getFn(); DoFnSignature signature = DoFnSignatures.signatureForDoFn(fn); return signature.processElement().isSplittable() && signature.isBoundedPerElement() == IsBounded.BOUNDED; } return false; }
@Override public boolean matches(AppliedPTransform<?, ?, ?> application) { PTransform<?, ?> transform = application.getTransform(); if (transform instanceof SplittableParDo.ProcessKeyedElements) { DoFn<?, ?> fn = ((SplittableParDo.ProcessKeyedElements) transform).getFn(); DoFnSignature signature = DoFnSignatures.signatureForDoFn(fn); return signature.processElement().isSplittable() && signature.isBoundedPerElement() == IsBounded.UNBOUNDED; } return false; }
@Override public PCollectionTuple expand( PCollection<KeyedWorkItem<byte[], KV<InputT, RestrictionT>>> input) { return ProcessKeyedElements.createPrimitiveOutputFor( input, original.getFn(), original.getMainOutputTag(), original.getAdditionalOutputTags(), original.getOutputTagsToCoders(), original.getInputWindowingStrategy()); } }