@Override public boolean matches(AppliedPTransform<?, ?, ?> application) { return application.getTransform().getClass().equals(Combine.GroupedValues.class) && ((Combine.GroupedValues<?, ?, ?>) application.getTransform()) .getSideInputs() .isEmpty(); }
private AppliedPTransform<?, ?, ?> getCurrentTransform(PTransform<?, ?> transform) { checkArgument( currentTransform != null && currentTransform.getTransform() == transform, "can only be called with current transform"); return currentTransform; }
public static PTransformMatcher createViewWithViewFn(final Class<? extends ViewFn> viewFnType) { return application -> { if (!(application.getTransform() instanceof CreatePCollectionView)) { return false; } CreatePCollectionView<?, ?> createView = (CreatePCollectionView<?, ?>) application.getTransform(); ViewFn<?, ?> viewFn = createView.getView().getViewFn(); return viewFn.getClass().equals(viewFnType); }; }
@Override public boolean appliesTo(PCollection<?> collection, DirectGraph graph) { return CONTAINS_UDF.contains( PTransformTranslation.urnForTransform(graph.getProducer(collection).getTransform())); } };
@Override public DoFnLifecycleManager load(final AppliedPTransform<?, ?, ?> application) { checkArgument( ProcessElements.class.isInstance(application.getTransform()), "No know extraction of the fn from " + application); final ProcessElements<InputT, OutputT, RestrictionT, PositionT> transform = (ProcessElements<InputT, OutputT, RestrictionT, PositionT>) application.getTransform(); return DoFnLifecycleManager.of(transform.newProcessFn(transform.getFn())); } },
public <InputT> TransformEvaluator<InputT> forApplication( AppliedPTransform<?, ?, ?> application, CommittedBundle<?> inputBundle) throws Exception { checkState( !finished.get(), "Tried to get an evaluator for a finished TransformEvaluatorRegistry"); String urn = PTransformTranslation.urnForTransform(application.getTransform()); TransformEvaluatorFactory factory = checkNotNull(factories.get(urn), "No evaluator for PTransform \"%s\"", urn); return factory.forApplication(application, inputBundle); }
@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<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)); }
public Map<TupleTag<?>, PValue> getOutputs(PTransform<?, ?> transform) { checkArgument(currentTransform != null, "can only be called with non-null currentTransform"); checkArgument( currentTransform.getTransform() == transform, "can only be called with current transform"); return currentTransform.getOutputs(); }
@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 PCollectionTuple expand(PCollection<InputT> input) { return input .apply("Materialize input", Reshuffle.viaRandomKey()) .apply("ParDo with stable input", appliedTransform.getTransform()); } });
@Override public PTransformReplacement<PBegin, PCollection<T>> getReplacementTransform( AppliedPTransform<PBegin, PCollection<T>, Read.Bounded<T>> transform) { return PTransformReplacement.of( transform.getPipeline().begin(), new StreamingBoundedRead<>(transform.getTransform())); }
@Override public FunctionSpec translate( AppliedPTransform<?, ?, GroupByKey<?, ?>> transform, SdkComponents components) { return FunctionSpec.newBuilder().setUrn(getUrn(transform.getTransform())).build(); } }
@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; }
@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 FunctionSpec translate( AppliedPTransform<?, ?, View.CreatePCollectionView<?, ?>> transform, SdkComponents components) { return FunctionSpec.newBuilder() .setUrn(getUrn(transform.getTransform())) .setPayload( ByteString.copyFrom( SerializableUtils.serializeToByteArray(transform.getTransform().getView()))) .build(); } }
@Override public boolean matches(AppliedPTransform<?, ?, ?> application) { PTransform<?, ?> transform = application.getTransform(); if (transform instanceof ParDo.MultiOutput) { DoFn<?, ?> fn = ((ParDo.MultiOutput<?, ?>) transform).getFn(); DoFnSignature signature = DoFnSignatures.signatureForDoFn(fn); return signature.usesState() || signature.usesTimers(); } return false; }
@Override public PTransformReplacement<PBegin, PCollection<T>> getReplacementTransform( AppliedPTransform<PBegin, PCollection<T>, Values<T>> transform) { return PTransformReplacement.of( transform.getPipeline().begin(), new PrimitiveCreate<>( transform.getTransform(), ((PCollection<T>) Iterables.getOnlyElement(transform.getOutputs().values())) .getCoder())); }
@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 FunctionSpec translate( AppliedPTransform<?, ?, Window.Assign<?>> transform, SdkComponents components) { return FunctionSpec.newBuilder() .setUrn("urn:beam:transform:window:v1") .setPayload( WindowIntoTranslation.toProto(transform.getTransform(), components).toByteString()) .build(); } }