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());
Map.Entry::getKey, e -> ((PCollection) e.getValue()).getCoder())); translateInputs( stepContext, context.getInput(transform), transform.getSideInputs(), context); translateOutputs(context.getOutputs(transform), stepContext); String ptransformId = stepContext, ptransformId, transform.getFn(), transform.getInputWindowingStrategy(), transform.getSideInputs(), transform.getElementCoder(), context, transform.getMainOutputTag(), outputCoders); translateCoder(transform.getRestrictionCoder(), context));
new ProcessKeyedElements<>( doFn, input.getCoder(),
@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(); }
public List<PCollectionView<?>> getSideInputs() { return original.getSideInputs(); }
public TupleTagList getAdditionalOutputTags() { return original.getAdditionalOutputTags(); }
public TupleTag<OutputT> getMainOutputTag() { return original.getMainOutputTag(); }
@Override public Map<String, SideInput> translateSideInputs(SdkComponents components) { return ParDoTranslation.translateSideInputs(pke.getSideInputs(), components); }
@Override public SdkFunctionSpec translateDoFn(SdkComponents newComponents) { return ParDoTranslation.translateDoFn(fn, pke.getMainOutputTag(), newComponents); }
@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<KV<byte[], KV<InputT, RestrictionT>>> input) { return createPrimitiveOutputFor( input, fn, mainOutputTag, additionalOutputTags, outputTagsToCoders, windowingStrategy); }
@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()); } }
public ProcessFn<InputT, OutputT, RestrictionT, PositionT> newProcessFn( DoFn<InputT, OutputT> fn) { return new ProcessFn<>( fn, original.getElementCoder(), original.getRestrictionCoder(), original.getInputWindowingStrategy()); }