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

Example 6 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();
        }
    };
}

Example 7 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 8 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 9 with DoFnSignature

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

the class BatchStatefulParDoOverrides method verifyFnIsStateful.

private static <InputT, OutputT> void verifyFnIsStateful(DoFn<InputT, OutputT> fn) {
    DoFnSignature signature = DoFnSignatures.getSignature(fn.getClass());
    // It is still correct to use this without state or timers, but a bad idea.
    // Since it is internal it should never be used wrong, so it is OK to crash.
    checkState(signature.usesState() || signature.usesTimers(), "%s used for %s that does not use state or timers.", BatchStatefulParDoOverrides.class.getSimpleName(), ParDo.class.getSimpleName());
}

Example 10 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);
        }
    }
}