Examples with StateId org.apache.beam.sdk.transforms.DoFn.StateId used on opensource projects

Example 1 with StateId

use of org.apache.beam.sdk.transforms.DoFn.StateId 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);
}