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

Example 46 with StateSpec

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

the class DoFnSignaturesSplittableDoFnTest method testSplittableProcessElementMustNotHaveUnsupportedParams.

@Test
// used via reflection
@SuppressWarnings("unused")
public void testSplittableProcessElementMustNotHaveUnsupportedParams() throws Exception {
    thrown.expect(IllegalArgumentException.class);
    thrown.expectMessage("Illegal parameter");
    thrown.expectMessage("ValueState");
    DoFn<Integer, String> doFn = new DoFn<Integer, String>() {

        @StateId("my-state-id")
        public final StateSpec<ValueState<String>> myStateSpec = StateSpecs.value(StringUtf8Coder.of());

        @ProcessElement
        public void method(DoFn<Integer, String>.ProcessContext context, SomeRestrictionTracker tracker, @StateId("my-state-id") ValueState<String> myState) {
        }
    };
    Method processElementMethod = null;
    for (Method method : doFn.getClass().getDeclaredMethods()) {
        if ("method".equals(method.getName())) {
            processElementMethod = method;
        }
    }
    checkState(processElementMethod != null);
    FnAnalysisContext context = FnAnalysisContext.create();
    context.addStateDeclaration(DoFnSignature.StateDeclaration.create("my-state-id", doFn.getClass().getField("myStateSpec"), new TypeDescriptor<ValueState<String>>() {
    }));
    DoFnSignatures.analyzeProcessElementMethod(errors(), new TypeDescriptor<DoFn<Integer, String>>() {
    }, processElementMethod, TypeDescriptor.of(Integer.class), TypeDescriptor.of(String.class), context);
}

Example 47 with StateSpec

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

the class DoFnSignaturesTest method testStateParameterAlwaysFetched.

@Test
public void testStateParameterAlwaysFetched() {
    thrown.expect(IllegalArgumentException.class);
    thrown.expectMessage("ReadableStates");
    DoFnSignature sig = DoFnSignatures.getSignature(new DoFn<KV<String, Integer>, Long>() {

        @StateId("my-id")
        private final StateSpec<MapState<Integer, Integer>> myfield = StateSpecs.map(VarIntCoder.of(), VarIntCoder.of());

        @ProcessElement
        public void myProcessElement(ProcessContext context, @AlwaysFetched @StateId("my-id") MapState<Integer, Integer> one) {
        }
    }.getClass());
    StateParameter stateParameter = (StateParameter) sig.processElement().extraParameters().get(1);
    assertTrue(stateParameter.alwaysFetched());
}