Examples with State org.apache.beam.sdk.state.State used on opensource projects

Example 1 with State

use of org.apache.beam.sdk.state.State in project beam by apache.

the class DoFnSignatures method analyzeStateDeclarations.

private static Map<String, DoFnSignature.StateDeclaration> analyzeStateDeclarations(ErrorReporter errors, Class<?> fnClazz) {
    Map<String, DoFnSignature.StateDeclaration> declarations = new HashMap<>();
    for (Field field : declaredFieldsWithAnnotation(DoFn.StateId.class, fnClazz, DoFn.class)) {
        // StateSpec fields may generally be private, but will be accessed via the signature
        field.setAccessible(true);
        String id = field.getAnnotation(DoFn.StateId.class).value();
        if (declarations.containsKey(id)) {
            errors.throwIllegalArgument("Duplicate %s \"%s\", used on both of [%s] and [%s]", DoFn.StateId.class.getSimpleName(), id, field.toString(), declarations.get(id).field().toString());
            continue;
        }
        Class<?> stateSpecRawType = field.getType();
        if (!(TypeDescriptor.of(stateSpecRawType).isSubtypeOf(TypeDescriptor.of(StateSpec.class)))) {
            errors.throwIllegalArgument("%s annotation on non-%s field [%s] that has class %s", DoFn.StateId.class.getSimpleName(), StateSpec.class.getSimpleName(), field.toString(), stateSpecRawType.getName());
            continue;
        }
        if (!Modifier.isFinal(field.getModifiers())) {
            errors.throwIllegalArgument("Non-final field %s annotated with %s. State declarations must be final.", field.toString(), DoFn.StateId.class.getSimpleName());
            continue;
        }
        Type stateSpecType = field.getGenericType();
        // A type descriptor for whatever type the @StateId-annotated class has, which
        // must be some subtype of StateSpec
        TypeDescriptor<? extends StateSpec<?>> stateSpecSubclassTypeDescriptor = (TypeDescriptor) TypeDescriptor.of(stateSpecType);
        // A type descriptor for StateSpec, with the generic type parameters filled
        // in according to the specialization of the subclass (or just straight params)
        TypeDescriptor<StateSpec<?>> stateSpecTypeDescriptor = (TypeDescriptor) stateSpecSubclassTypeDescriptor.getSupertype(StateSpec.class);
        // The type of the state, which may still have free type variables from the
        // context
        Type unresolvedStateType = ((ParameterizedType) stateSpecTypeDescriptor.getType()).getActualTypeArguments()[0];
        // By static typing this is already a well-formed State subclass
        TypeDescriptor<? extends State> stateType = (TypeDescriptor<? extends State>) TypeDescriptor.of(fnClazz).resolveType(unresolvedStateType);
        declarations.put(id, DoFnSignature.StateDeclaration.create(id, field, stateType));
    }
    return ImmutableMap.copyOf(declarations);
}

Example 2 with State

use of org.apache.beam.sdk.state.State in project beam by apache.

the class StateTable method get.

/**
   * Gets the {@link State} in the specified {@link StateNamespace} with the specified {@link
   * StateTag}, binding it using the {@link #binderForNamespace} if it is not
   * already present in this {@link StateTable}.
   */
public <StateT extends State> StateT get(StateNamespace namespace, StateTag<StateT> tag, StateContext<?> c) {
    State storage = stateTable.get(namespace, tag);
    if (storage != null) {
        @SuppressWarnings("unchecked") StateT typedStorage = (StateT) storage;
        return typedStorage;
    }
    StateT typedStorage = tag.bind(binderForNamespace(namespace, c));
    stateTable.put(namespace, tag, typedStorage);
    return typedStorage;
}

Example 3 with State

use of org.apache.beam.sdk.state.State 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;
    }
}