Examples with DoFnSignature org.apache.beam.sdk.transforms.reflect.DoFnSignature used on opensource projects

Example 16 with DoFnSignature

use of org.apache.beam.sdk.transforms.reflect.DoFnSignature in project beam by apache.

the class ParDoTranslation method toProto.

public static ParDoPayload toProto(ParDo.MultiOutput<?, ?> parDo, SdkComponents components) throws IOException {
    DoFn<?, ?> doFn = parDo.getFn();
    DoFnSignature signature = DoFnSignatures.getSignature(doFn.getClass());
    Map<String, StateDeclaration> states = signature.stateDeclarations();
    Map<String, TimerDeclaration> timers = signature.timerDeclarations();
    List<Parameter> parameters = signature.processElement().extraParameters();
    ParDoPayload.Builder builder = ParDoPayload.newBuilder();
    builder.setDoFn(toProto(parDo.getFn(), parDo.getMainOutputTag()));
    for (PCollectionView<?> sideInput : parDo.getSideInputs()) {
        builder.putSideInputs(sideInput.getTagInternal().getId(), toProto(sideInput));
    }
    for (Parameter parameter : parameters) {
        Optional<RunnerApi.Parameter> protoParameter = toProto(parameter);
        if (protoParameter.isPresent()) {
            builder.addParameters(protoParameter.get());
        }
    }
    for (Map.Entry<String, StateDeclaration> state : states.entrySet()) {
        RunnerApi.StateSpec spec = toProto(getStateSpecOrCrash(state.getValue(), doFn), components);
        builder.putStateSpecs(state.getKey(), spec);
    }
    for (Map.Entry<String, TimerDeclaration> timer : timers.entrySet()) {
        RunnerApi.TimerSpec spec = toProto(getTimerSpecOrCrash(timer.getValue(), doFn));
        builder.putTimerSpecs(timer.getKey(), spec);
    }
    return builder.build();
}

Example 17 with DoFnSignature

use of org.apache.beam.sdk.transforms.reflect.DoFnSignature in project beam by apache.

the class DataflowPipelineTranslator method translateFn.

private static void translateFn(StepTranslationContext stepContext, DoFn fn, WindowingStrategy windowingStrategy, Iterable<PCollectionView<?>> sideInputs, Coder inputCoder, TranslationContext context, long mainOutput, Map<Long, TupleTag<?>> outputMap) {
    DoFnSignature signature = DoFnSignatures.getSignature(fn.getClass());
    if (signature.processElement().isSplittable()) {
        throw new UnsupportedOperationException(String.format("%s does not currently support splittable DoFn: %s", DataflowRunner.class.getSimpleName(), fn));
    }
    stepContext.addInput(PropertyNames.USER_FN, fn.getClass().getName());
    stepContext.addInput(PropertyNames.SERIALIZED_FN, byteArrayToJsonString(serializeToByteArray(DoFnInfo.forFn(fn, windowingStrategy, sideInputs, inputCoder, mainOutput, outputMap))));
    // in streaming but does not work in batch
    if (context.getPipelineOptions().isStreaming() && (signature.usesState() || signature.usesTimers())) {
        stepContext.addInput(PropertyNames.USES_KEYED_STATE, "true");
    }
}

Example 18 with DoFnSignature

use of org.apache.beam.sdk.transforms.reflect.DoFnSignature in project beam by apache.

the class ParDo method finishSpecifyingStateSpecs.

private static void finishSpecifyingStateSpecs(DoFn<?, ?> fn, CoderRegistry coderRegistry, Coder<?> inputCoder) {
    DoFnSignature signature = DoFnSignatures.getSignature(fn.getClass());
    Map<String, DoFnSignature.StateDeclaration> stateDeclarations = signature.stateDeclarations();
    for (DoFnSignature.StateDeclaration stateDeclaration : stateDeclarations.values()) {
        try {
            StateSpec<?> stateSpec = (StateSpec<?>) stateDeclaration.field().get(fn);
            stateSpec.offerCoders(codersForStateSpecTypes(stateDeclaration, coderRegistry, inputCoder));
            stateSpec.finishSpecifying();
        } catch (IllegalAccessException e) {
            throw new RuntimeException(e);
        }
    }
}

Example 19 with DoFnSignature

use of org.apache.beam.sdk.transforms.reflect.DoFnSignature in project beam by apache.

the class ParDoMultiOverrideFactory method getReplacementTransform.

@SuppressWarnings("unchecked")
private PTransform<PCollection<? extends InputT>, PCollectionTuple> getReplacementTransform(MultiOutput<InputT, OutputT> transform) {
    DoFn<InputT, OutputT> fn = transform.getFn();
    DoFnSignature signature = DoFnSignatures.getSignature(fn.getClass());
    if (signature.processElement().isSplittable()) {
        return new SplittableParDo(transform);
    } else if (signature.stateDeclarations().size() > 0 || signature.timerDeclarations().size() > 0) {
        // Based on the fact that the signature is stateful, DoFnSignatures ensures
        // that it is also keyed
        MultiOutput<KV<?, ?>, OutputT> keyedTransform = (MultiOutput<KV<?, ?>, OutputT>) transform;
        return new GbkThenStatefulParDo(keyedTransform);
    } else {
        return transform;
    }
}

Example 20 with DoFnSignature

use of org.apache.beam.sdk.transforms.reflect.DoFnSignature in project beam by apache.

the class PTransformMatchers method stateOrTimerParDoMulti.

/**
 * A {@link PTransformMatcher} that matches a {@link ParDo.MultiOutput} containing a {@link DoFn}
 * that uses state or timers, as specified by {@link DoFnSignature#usesState()} and {@link
 * DoFnSignature#usesTimers()}.
 */
public static PTransformMatcher stateOrTimerParDoMulti() {
    return new PTransformMatcher() {

        @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 String toString() {
            return MoreObjects.toStringHelper("StateOrTimerParDoMultiMatcher").toString();
        }
    };
}