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

Example 1 with TimerDeclaration

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

the class DoFnSignatures method analyzeTimerDeclarations.

private static ImmutableMap<String, TimerDeclaration> analyzeTimerDeclarations(ErrorReporter errors, Class<?> fnClazz) {
    Map<String, DoFnSignature.TimerDeclaration> declarations = new HashMap<>();
    for (Field field : declaredFieldsWithAnnotation(DoFn.TimerId.class, fnClazz, DoFn.class)) {
        // TimerSpec fields may generally be private, but will be accessed via the signature
        field.setAccessible(true);
        String id = field.getAnnotation(DoFn.TimerId.class).value();
        validateTimerField(errors, declarations, id, field);
        declarations.put(id, DoFnSignature.TimerDeclaration.create(id, field));
    }
    return ImmutableMap.copyOf(declarations);
}

Example 2 with TimerDeclaration

use of org.apache.beam.sdk.transforms.reflect.DoFnSignature.TimerDeclaration 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 3 with TimerDeclaration

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

the class DoFnSignatures method analyzeExtraParameter.

private static Parameter analyzeExtraParameter(ErrorReporter methodErrors, FnAnalysisContext fnContext, MethodAnalysisContext methodContext, TypeDescriptor<? extends DoFn<?, ?>> fnClass, ParameterDescription param, TypeDescriptor<?> inputT, TypeDescriptor<?> outputT) {
    TypeDescriptor<?> expectedProcessContextT = doFnProcessContextTypeOf(inputT, outputT);
    TypeDescriptor<?> expectedOnTimerContextT = doFnOnTimerContextTypeOf(inputT, outputT);
    TypeDescriptor<?> paramT = param.getType();
    Class<?> rawType = paramT.getRawType();
    ErrorReporter paramErrors = methodErrors.forParameter(param);
    if (rawType.equals(DoFn.ProcessContext.class)) {
        paramErrors.checkArgument(paramT.equals(expectedProcessContextT), "ProcessContext argument must have type %s", formatType(expectedProcessContextT));
        return Parameter.processContext();
    } else if (rawType.equals(DoFn.OnTimerContext.class)) {
        paramErrors.checkArgument(paramT.equals(expectedOnTimerContextT), "OnTimerContext argument must have type %s", formatType(expectedOnTimerContextT));
        return Parameter.onTimerContext();
    } else if (BoundedWindow.class.isAssignableFrom(rawType)) {
        methodErrors.checkArgument(!methodContext.hasWindowParameter(), "Multiple %s parameters", BoundedWindow.class.getSimpleName());
        return Parameter.boundedWindow((TypeDescriptor<? extends BoundedWindow>) paramT);
    } else if (RestrictionTracker.class.isAssignableFrom(rawType)) {
        methodErrors.checkArgument(!methodContext.hasRestrictionTrackerParameter(), "Multiple %s parameters", RestrictionTracker.class.getSimpleName());
        return Parameter.restrictionTracker(paramT);
    } else if (rawType.equals(Timer.class)) {
        // m.getParameters() is not available until Java 8
        String id = getTimerId(param.getAnnotations());
        paramErrors.checkArgument(id != null, "%s missing %s annotation", Timer.class.getSimpleName(), TimerId.class.getSimpleName());
        paramErrors.checkArgument(!methodContext.getTimerParameters().containsKey(id), "duplicate %s: \"%s\"", TimerId.class.getSimpleName(), id);
        TimerDeclaration timerDecl = fnContext.getTimerDeclarations().get(id);
        paramErrors.checkArgument(timerDecl != null, "reference to undeclared %s: \"%s\"", TimerId.class.getSimpleName(), id);
        paramErrors.checkArgument(timerDecl.field().getDeclaringClass().equals(param.getMethod().getDeclaringClass()), "%s %s declared in a different class %s." + " Timers may be referenced only in the lexical scope where they are declared.", TimerId.class.getSimpleName(), id, timerDecl.field().getDeclaringClass().getName());
        return Parameter.timerParameter(timerDecl);
    } else if (State.class.isAssignableFrom(rawType)) {
        // m.getParameters() is not available until Java 8
        String id = getStateId(param.getAnnotations());
        paramErrors.checkArgument(id != null, "missing %s annotation", DoFn.StateId.class.getSimpleName());
        paramErrors.checkArgument(!methodContext.getStateParameters().containsKey(id), "duplicate %s: \"%s\"", DoFn.StateId.class.getSimpleName(), id);
        // By static typing this is already a well-formed State subclass
        TypeDescriptor<? extends State> stateType = (TypeDescriptor<? extends State>) param.getType();
        StateDeclaration stateDecl = fnContext.getStateDeclarations().get(id);
        paramErrors.checkArgument(stateDecl != null, "reference to undeclared %s: \"%s\"", DoFn.StateId.class.getSimpleName(), id);
        paramErrors.checkArgument(stateDecl.stateType().equals(stateType), "reference to %s %s with different type %s", StateId.class.getSimpleName(), id, formatType(stateDecl.stateType()));
        paramErrors.checkArgument(stateDecl.field().getDeclaringClass().equals(param.getMethod().getDeclaringClass()), "%s %s declared in a different class %s." + " State may be referenced only in the class where it is declared.", StateId.class.getSimpleName(), id, stateDecl.field().getDeclaringClass().getName());
        return Parameter.stateParameter(stateDecl);
    } else {
        List<String> allowedParamTypes = Arrays.asList(formatType(new TypeDescriptor<BoundedWindow>() {
        }), formatType(new TypeDescriptor<RestrictionTracker<?>>() {
        }));
        paramErrors.throwIllegalArgument("%s is not a valid context parameter. Should be one of %s", formatType(paramT), allowedParamTypes);
        // Unreachable
        return null;
    }
}

Example 4 with TimerDeclaration

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

the class DoFnSignatures method parseSignature.

/** Analyzes a given {@link DoFn} class and extracts its {@link DoFnSignature}. */
private static DoFnSignature parseSignature(Class<? extends DoFn<?, ?>> fnClass) {
    DoFnSignature.Builder signatureBuilder = DoFnSignature.builder();
    ErrorReporter errors = new ErrorReporter(null, fnClass.getName());
    errors.checkArgument(DoFn.class.isAssignableFrom(fnClass), "Must be subtype of DoFn");
    signatureBuilder.setFnClass(fnClass);
    TypeDescriptor<? extends DoFn<?, ?>> fnT = TypeDescriptor.of(fnClass);
    // Extract the input and output type, and whether the fn is bounded.
    TypeDescriptor<?> inputT = null;
    TypeDescriptor<?> outputT = null;
    for (TypeDescriptor<?> supertype : fnT.getTypes()) {
        if (!supertype.getRawType().equals(DoFn.class)) {
            continue;
        }
        Type[] args = ((ParameterizedType) supertype.getType()).getActualTypeArguments();
        inputT = TypeDescriptor.of(args[0]);
        outputT = TypeDescriptor.of(args[1]);
    }
    errors.checkNotNull(inputT, "Unable to determine input type");
    // Find the state and timer declarations in advance of validating
    // method parameter lists
    FnAnalysisContext fnContext = FnAnalysisContext.create();
    fnContext.addStateDeclarations(analyzeStateDeclarations(errors, fnClass).values());
    fnContext.addTimerDeclarations(analyzeTimerDeclarations(errors, fnClass).values());
    Method processElementMethod = findAnnotatedMethod(errors, DoFn.ProcessElement.class, fnClass, true);
    Method startBundleMethod = findAnnotatedMethod(errors, DoFn.StartBundle.class, fnClass, false);
    Method finishBundleMethod = findAnnotatedMethod(errors, DoFn.FinishBundle.class, fnClass, false);
    Method setupMethod = findAnnotatedMethod(errors, DoFn.Setup.class, fnClass, false);
    Method teardownMethod = findAnnotatedMethod(errors, DoFn.Teardown.class, fnClass, false);
    Method getInitialRestrictionMethod = findAnnotatedMethod(errors, DoFn.GetInitialRestriction.class, fnClass, false);
    Method splitRestrictionMethod = findAnnotatedMethod(errors, DoFn.SplitRestriction.class, fnClass, false);
    Method getRestrictionCoderMethod = findAnnotatedMethod(errors, DoFn.GetRestrictionCoder.class, fnClass, false);
    Method newTrackerMethod = findAnnotatedMethod(errors, DoFn.NewTracker.class, fnClass, false);
    Collection<Method> onTimerMethods = declaredMethodsWithAnnotation(DoFn.OnTimer.class, fnClass, DoFn.class);
    HashMap<String, DoFnSignature.OnTimerMethod> onTimerMethodMap = Maps.newHashMapWithExpectedSize(onTimerMethods.size());
    for (Method onTimerMethod : onTimerMethods) {
        String id = onTimerMethod.getAnnotation(DoFn.OnTimer.class).value();
        errors.checkArgument(fnContext.getTimerDeclarations().containsKey(id), "Callback %s is for undeclared timer %s", onTimerMethod, id);
        TimerDeclaration timerDecl = fnContext.getTimerDeclarations().get(id);
        errors.checkArgument(timerDecl.field().getDeclaringClass().equals(onTimerMethod.getDeclaringClass()), "Callback %s is for timer %s declared in a different class %s." + " Timer callbacks must be declared in the same lexical scope as their timer", onTimerMethod, id, timerDecl.field().getDeclaringClass().getCanonicalName());
        onTimerMethodMap.put(id, analyzeOnTimerMethod(errors, fnT, onTimerMethod, id, inputT, outputT, fnContext));
    }
    signatureBuilder.setOnTimerMethods(onTimerMethodMap);
    // those timers used in methods, once method parameter lists support timers.
    for (TimerDeclaration decl : fnContext.getTimerDeclarations().values()) {
        errors.checkArgument(onTimerMethodMap.containsKey(decl.id()), "No callback registered via %s for timer %s", DoFn.OnTimer.class.getSimpleName(), decl.id());
    }
    ErrorReporter processElementErrors = errors.forMethod(DoFn.ProcessElement.class, processElementMethod);
    DoFnSignature.ProcessElementMethod processElement = analyzeProcessElementMethod(processElementErrors, fnT, processElementMethod, inputT, outputT, fnContext);
    signatureBuilder.setProcessElement(processElement);
    if (startBundleMethod != null) {
        ErrorReporter startBundleErrors = errors.forMethod(DoFn.StartBundle.class, startBundleMethod);
        signatureBuilder.setStartBundle(analyzeStartBundleMethod(startBundleErrors, fnT, startBundleMethod, inputT, outputT));
    }
    if (finishBundleMethod != null) {
        ErrorReporter finishBundleErrors = errors.forMethod(DoFn.FinishBundle.class, finishBundleMethod);
        signatureBuilder.setFinishBundle(analyzeFinishBundleMethod(finishBundleErrors, fnT, finishBundleMethod, inputT, outputT));
    }
    if (setupMethod != null) {
        signatureBuilder.setSetup(analyzeLifecycleMethod(errors.forMethod(DoFn.Setup.class, setupMethod), setupMethod));
    }
    if (teardownMethod != null) {
        signatureBuilder.setTeardown(analyzeLifecycleMethod(errors.forMethod(DoFn.Teardown.class, teardownMethod), teardownMethod));
    }
    ErrorReporter getInitialRestrictionErrors;
    if (getInitialRestrictionMethod != null) {
        getInitialRestrictionErrors = errors.forMethod(DoFn.GetInitialRestriction.class, getInitialRestrictionMethod);
        signatureBuilder.setGetInitialRestriction(analyzeGetInitialRestrictionMethod(getInitialRestrictionErrors, fnT, getInitialRestrictionMethod, inputT));
    }
    if (splitRestrictionMethod != null) {
        ErrorReporter splitRestrictionErrors = errors.forMethod(DoFn.SplitRestriction.class, splitRestrictionMethod);
        signatureBuilder.setSplitRestriction(analyzeSplitRestrictionMethod(splitRestrictionErrors, fnT, splitRestrictionMethod, inputT));
    }
    if (getRestrictionCoderMethod != null) {
        ErrorReporter getRestrictionCoderErrors = errors.forMethod(DoFn.GetRestrictionCoder.class, getRestrictionCoderMethod);
        signatureBuilder.setGetRestrictionCoder(analyzeGetRestrictionCoderMethod(getRestrictionCoderErrors, fnT, getRestrictionCoderMethod));
    }
    if (newTrackerMethod != null) {
        ErrorReporter newTrackerErrors = errors.forMethod(DoFn.NewTracker.class, newTrackerMethod);
        signatureBuilder.setNewTracker(analyzeNewTrackerMethod(newTrackerErrors, fnT, newTrackerMethod));
    }
    signatureBuilder.setIsBoundedPerElement(inferBoundedness(fnT, processElement, errors));
    signatureBuilder.setStateDeclarations(fnContext.getStateDeclarations());
    signatureBuilder.setTimerDeclarations(fnContext.getTimerDeclarations());
    DoFnSignature signature = signatureBuilder.build();
    // Additional validation for splittable DoFn's.
    if (processElement.isSplittable()) {
        verifySplittableMethods(signature, errors);
    } else {
        verifyUnsplittableMethods(errors, signature);
    }
    return signature;
}