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

Example 11 with ValueState

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

the class ParDoTest method testValueStateSameId.

/**
   * Tests that there is no state bleeding between adjacent stateful {@link ParDo} transforms,
   * which may (or may not) be executed in similar contexts after runner optimizations.
   */
@Test
@Category({ ValidatesRunner.class, UsesStatefulParDo.class })
public void testValueStateSameId() {
    final String stateId = "foo";
    DoFn<KV<String, Integer>, KV<String, Integer>> fn = new DoFn<KV<String, Integer>, KV<String, Integer>>() {

        @StateId(stateId)
        private final StateSpec<ValueState<Integer>> intState = StateSpecs.value(VarIntCoder.of());

        @ProcessElement
        public void processElement(ProcessContext c, @StateId(stateId) ValueState<Integer> state) {
            Integer currentValue = MoreObjects.firstNonNull(state.read(), 0);
            c.output(KV.of("sizzle", currentValue));
            state.write(currentValue + 1);
        }
    };
    DoFn<KV<String, Integer>, Integer> fn2 = new DoFn<KV<String, Integer>, Integer>() {

        @StateId(stateId)
        private final StateSpec<ValueState<Integer>> intState = StateSpecs.value(VarIntCoder.of());

        @ProcessElement
        public void processElement(ProcessContext c, @StateId(stateId) ValueState<Integer> state) {
            Integer currentValue = MoreObjects.firstNonNull(state.read(), 13);
            c.output(currentValue);
            state.write(currentValue + 13);
        }
    };
    PCollection<KV<String, Integer>> intermediate = pipeline.apply(Create.of(KV.of("hello", 42), KV.of("hello", 97), KV.of("hello", 84))).apply("First stateful ParDo", ParDo.of(fn));
    PCollection<Integer> output = intermediate.apply("Second stateful ParDo", ParDo.of(fn2));
    PAssert.that(intermediate).containsInAnyOrder(KV.of("sizzle", 0), KV.of("sizzle", 1), KV.of("sizzle", 2));
    PAssert.that(output).containsInAnyOrder(13, 26, 39);
    pipeline.run();
}

Example 12 with ValueState

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

the class ParDoTest method testValueStateSimple.

@Test
@Category({ ValidatesRunner.class, UsesStatefulParDo.class })
public void testValueStateSimple() {
    final String stateId = "foo";
    DoFn<KV<String, Integer>, Integer> fn = new DoFn<KV<String, Integer>, Integer>() {

        @StateId(stateId)
        private final StateSpec<ValueState<Integer>> intState = StateSpecs.value(VarIntCoder.of());

        @ProcessElement
        public void processElement(ProcessContext c, @StateId(stateId) ValueState<Integer> state) {
            Integer currentValue = MoreObjects.firstNonNull(state.read(), 0);
            c.output(currentValue);
            state.write(currentValue + 1);
        }
    };
    PCollection<Integer> output = pipeline.apply(Create.of(KV.of("hello", 42), KV.of("hello", 97), KV.of("hello", 84))).apply(ParDo.of(fn));
    PAssert.that(output).containsInAnyOrder(0, 1, 2);
    pipeline.run();
}

Example 13 with ValueState

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

the class ParDoTest method testValueStateCoderInferenceFromInputCoder.

@Test
@Category({ ValidatesRunner.class, UsesStatefulParDo.class })
public void testValueStateCoderInferenceFromInputCoder() {
    final String stateId = "foo";
    MyIntegerCoder myIntegerCoder = MyIntegerCoder.of();
    DoFn<KV<String, MyInteger>, MyInteger> fn = new DoFn<KV<String, MyInteger>, MyInteger>() {

        @StateId(stateId)
        private final StateSpec<ValueState<MyInteger>> intState = StateSpecs.value();

        @ProcessElement
        public void processElement(ProcessContext c, @StateId(stateId) ValueState<MyInteger> state) {
            MyInteger currentValue = MoreObjects.firstNonNull(state.read(), new MyInteger(0));
            c.output(currentValue);
            state.write(new MyInteger(currentValue.getValue() + 1));
        }
    };
    pipeline.apply(Create.of(KV.of("hello", new MyInteger(42)), KV.of("hello", new MyInteger(97)), KV.of("hello", new MyInteger(84))).withCoder(KvCoder.of(StringUtf8Coder.of(), myIntegerCoder))).apply(ParDo.of(fn)).setCoder(myIntegerCoder);
    pipeline.run();
}

Example 14 with ValueState

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

the class CopyOnAccessInMemoryStateInternalsTest method testGetWithPresentInUnderlying.

/**
   * Tests that retrieving state with an underlying StateInternals with an existing value returns
   * a value that initially has equal value to the provided state but can be modified without
   * modifying the existing state.
   */
@Test
public void testGetWithPresentInUnderlying() {
    CopyOnAccessInMemoryStateInternals<String> underlying = CopyOnAccessInMemoryStateInternals.withUnderlying(key, null);
    StateNamespace namespace = new StateNamespaceForTest("foo");
    StateTag<ValueState<String>> valueTag = StateTags.value("foo", StringUtf8Coder.of());
    ValueState<String> underlyingValue = underlying.state(namespace, valueTag);
    assertThat(underlyingValue.read(), nullValue(String.class));
    underlyingValue.write("bar");
    assertThat(underlyingValue.read(), equalTo("bar"));
    CopyOnAccessInMemoryStateInternals<String> internals = CopyOnAccessInMemoryStateInternals.withUnderlying(key, underlying);
    ValueState<String> copyOnAccessState = internals.state(namespace, valueTag);
    assertThat(copyOnAccessState.read(), equalTo("bar"));
    copyOnAccessState.write("baz");
    assertThat(copyOnAccessState.read(), equalTo("baz"));
    assertThat(underlyingValue.read(), equalTo("bar"));
    ValueState<String> reReadUnderlyingValue = underlying.state(namespace, valueTag);
    assertThat(underlyingValue.read(), equalTo(reReadUnderlyingValue.read()));
}

Example 15 with ValueState

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

the class ParDoTest method testValueStateCoderInferenceFailure.

@Test
@Category({ ValidatesRunner.class, UsesStatefulParDo.class })
public void testValueStateCoderInferenceFailure() throws Exception {
    final String stateId = "foo";
    MyIntegerCoder myIntegerCoder = MyIntegerCoder.of();
    DoFn<KV<String, Integer>, MyInteger> fn = new DoFn<KV<String, Integer>, MyInteger>() {

        @StateId(stateId)
        private final StateSpec<ValueState<MyInteger>> intState = StateSpecs.value();

        @ProcessElement
        public void processElement(ProcessContext c, @StateId(stateId) ValueState<MyInteger> state) {
            MyInteger currentValue = MoreObjects.firstNonNull(state.read(), new MyInteger(0));
            c.output(currentValue);
            state.write(new MyInteger(currentValue.getValue() + 1));
        }
    };
    thrown.expect(RuntimeException.class);
    thrown.expectMessage("Unable to infer a coder for ValueState and no Coder was specified.");
    pipeline.apply(Create.of(KV.of("hello", 42), KV.of("hello", 97), KV.of("hello", 84))).apply(ParDo.of(fn)).setCoder(myIntegerCoder);
    pipeline.run();
}