@Override protected String getKindString() { return String.format("ParDo(%s)", NameUtils.approximateSimpleName(getFn())); }
@Override public boolean matches(AppliedPTransform<?, ?, ?> application) { DoFn<?, ?> fn; if (application.getTransform() instanceof ParDo.SingleOutput) { fn = ((ParDo.SingleOutput) application.getTransform()).getFn(); } else if (application.getTransform() instanceof ParDo.MultiOutput) { fn = ((ParDo.MultiOutput) application.getTransform()).getFn(); } else { return false; } return fnType.equals(fn.getClass()); }
@Test public void getReplacementTransformGetFn() { DoFn<Integer, Long> originalFn = new ToLongFn(); ParDo.SingleOutput<Integer, Long> originalTransform = ParDo.of(originalFn); PCollection<? extends Integer> input = pipeline.apply(Create.of(1, 2, 3)); AppliedPTransform< PCollection<? extends Integer>, PCollection<Long>, ParDo.SingleOutput<Integer, Long>> application = AppliedPTransform.of( "original", input.expand(), input.apply(originalTransform).expand(), originalTransform, pipeline); PTransformReplacement<PCollection<? extends Integer>, PCollection<Long>> replacementTransform = factory.getReplacementTransform(application); ParDoSingle<Integer, Long> parDoSingle = (ParDoSingle<Integer, Long>) replacementTransform.getTransform(); assertThat(parDoSingle.getFn(), equalTo(originalTransform.getFn())); assertThat(parDoSingle.getFn(), equalTo(originalFn)); }
@Override public PCollection<OutputT> expand(PCollection<? extends InputT> input) { SchemaRegistry schemaRegistry = input.getPipeline().getSchemaRegistry(); CoderRegistry registry = input.getPipeline().getCoderRegistry(); finishSpecifyingStateSpecs(fn, registry, input.getCoder()); TupleTag<OutputT> mainOutput = new TupleTag<>(MAIN_OUTPUT_TAG); PCollection<OutputT> res = input.apply(withOutputTags(mainOutput, TupleTagList.empty())).get(mainOutput); try { res.setSchema( schemaRegistry.getSchema(getFn().getOutputTypeDescriptor()), schemaRegistry.getToRowFunction(getFn().getOutputTypeDescriptor()), schemaRegistry.getFromRowFunction(getFn().getOutputTypeDescriptor())); } catch (NoSuchSchemaException e) { try { res.setCoder( registry.getCoder( getFn().getOutputTypeDescriptor(), getFn().getInputTypeDescriptor(), ((PCollection<InputT>) input).getCoder())); } catch (CannotProvideCoderException e2) { // Ignore and leave coder unset. } } return res; }
@Override public PCollection<OutputT> expand(PCollection<KV<K, InputT>> input) { DoFn<KV<K, InputT>, OutputT> fn = originalParDo.getFn(); verifyFnIsStateful(fn); DataflowRunner.verifyStateSupported(fn); DataflowRunner.verifyStateSupportForWindowingStrategy(input.getWindowingStrategy()); if (isFnApi) { return input .apply(GroupByKey.create()) .apply(ParDo.of(new ExpandGbkFn<>())) .apply(originalParDo); } PTransform< PCollection<? extends KV<K, Iterable<KV<Instant, WindowedValue<KV<K, InputT>>>>>>, PCollection<OutputT>> statefulParDo = ParDo.of(new BatchStatefulDoFn<>(fn)).withSideInputs(originalParDo.getSideInputs()); return input.apply(new GbkBeforeStatefulParDo<>()).apply(statefulParDo); } }
@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; } } }
public DoFn<InputT, OutputT> getFn() { return original.getFn(); }
@Override public boolean matches(AppliedPTransform<?, ?, ?> application) { PTransform<?, ?> transform = application.getTransform(); if (transform instanceof ParDo.SingleOutput) { DoFn<?, ?> fn = ((ParDo.SingleOutput<?, ?>) transform).getFn(); DoFnSignature signature = DoFnSignatures.signatureForDoFn(fn); return signature.processElement().isSplittable(); } return false; }
@Override public boolean matches(AppliedPTransform<?, ?, ?> application) { PTransform<?, ?> transform = application.getTransform(); if (transform instanceof ParDo.SingleOutput) { DoFn<?, ?> fn = ((ParDo.SingleOutput<?, ?>) transform).getFn(); DoFnSignature signature = DoFnSignatures.signatureForDoFn(fn); return signature.usesState() || signature.usesTimers(); } return false; }
@Override public boolean matches(AppliedPTransform<?, ?, ?> application) { PTransform<?, ?> transform = application.getTransform(); if (transform instanceof ParDo.SingleOutput) { DoFn<?, ?> fn = ((ParDo.SingleOutput<?, ?>) transform).getFn(); DoFnSignature signature = DoFnSignatures.signatureForDoFn(fn); return signature.processElement().requiresStableInput(); } return false; }