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

Example 1 with SetState

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

the class ParDoTest method testSetStateCoderInferenceFailure.

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

        @StateId(stateId)
        private final StateSpec<SetState<MyInteger>> setState = StateSpecs.set();

        @StateId(countStateId)
        private final StateSpec<CombiningState<Integer, int[], Integer>> countState = StateSpecs.combiningFromInputInternal(VarIntCoder.of(), Sum.ofIntegers());

        @ProcessElement
        public void processElement(ProcessContext c, @StateId(stateId) SetState<MyInteger> state, @StateId(countStateId) CombiningState<Integer, int[], Integer> count) {
            state.add(new MyInteger(c.element().getValue()));
            count.add(1);
            if (count.read() >= 4) {
                Set<MyInteger> set = Sets.newHashSet(state.read());
                c.output(set);
            }
        }
    };
    thrown.expect(RuntimeException.class);
    thrown.expectMessage("Unable to infer a coder for SetState and no Coder was specified.");
    pipeline.apply(Create.of(KV.of("hello", 97), KV.of("hello", 42), KV.of("hello", 42), KV.of("hello", 12))).apply(ParDo.of(fn)).setCoder(SetCoder.of(myIntegerCoder));
    pipeline.run();
}

Example 2 with SetState

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

the class ParDoTest method testSetStateCoderInference.

@Test
@Category({ ValidatesRunner.class, UsesStatefulParDo.class, UsesSetState.class })
public void testSetStateCoderInference() {
    final String stateId = "foo";
    final String countStateId = "count";
    Coder<MyInteger> myIntegerCoder = MyIntegerCoder.of();
    pipeline.getCoderRegistry().registerCoderForClass(MyInteger.class, myIntegerCoder);
    DoFn<KV<String, Integer>, Set<MyInteger>> fn = new DoFn<KV<String, Integer>, Set<MyInteger>>() {

        @StateId(stateId)
        private final StateSpec<SetState<MyInteger>> setState = StateSpecs.set();

        @StateId(countStateId)
        private final StateSpec<CombiningState<Integer, int[], Integer>> countState = StateSpecs.combiningFromInputInternal(VarIntCoder.of(), Sum.ofIntegers());

        @ProcessElement
        public void processElement(ProcessContext c, @StateId(stateId) SetState<MyInteger> state, @StateId(countStateId) CombiningState<Integer, int[], Integer> count) {
            state.add(new MyInteger(c.element().getValue()));
            count.add(1);
            if (count.read() >= 4) {
                Set<MyInteger> set = Sets.newHashSet(state.read());
                c.output(set);
            }
        }
    };
    PCollection<Set<MyInteger>> output = pipeline.apply(Create.of(KV.of("hello", 97), KV.of("hello", 42), KV.of("hello", 42), KV.of("hello", 12))).apply(ParDo.of(fn)).setCoder(SetCoder.of(myIntegerCoder));
    PAssert.that(output).containsInAnyOrder(Sets.newHashSet(new MyInteger(97), new MyInteger(42), new MyInteger(12)));
    pipeline.run();
}

Example 3 with SetState

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

the class ParDoTest method testSetState.

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

        @StateId(stateId)
        private final StateSpec<SetState<Integer>> setState = StateSpecs.set(VarIntCoder.of());

        @StateId(countStateId)
        private final StateSpec<CombiningState<Integer, int[], Integer>> countState = StateSpecs.combiningFromInputInternal(VarIntCoder.of(), Sum.ofIntegers());

        @ProcessElement
        public void processElement(ProcessContext c, @StateId(stateId) SetState<Integer> state, @StateId(countStateId) CombiningState<Integer, int[], Integer> count) {
            state.add(c.element().getValue());
            count.add(1);
            if (count.read() >= 4) {
                Set<Integer> set = Sets.newHashSet(state.read());
                c.output(set);
            }
        }
    };
    PCollection<Set<Integer>> output = pipeline.apply(Create.of(KV.of("hello", 97), KV.of("hello", 42), KV.of("hello", 42), KV.of("hello", 12))).apply(ParDo.of(fn));
    PAssert.that(output).containsInAnyOrder(Sets.newHashSet(97, 42, 12));
    pipeline.run();
}

Example 4 with SetState

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

the class CopyOnAccessInMemoryStateInternalsTest method testSetStateWithUnderlying.

@Test
public void testSetStateWithUnderlying() {
    CopyOnAccessInMemoryStateInternals<String> underlying = CopyOnAccessInMemoryStateInternals.withUnderlying(key, null);
    StateNamespace namespace = new StateNamespaceForTest("foo");
    StateTag<SetState<Integer>> valueTag = StateTags.set("foo", VarIntCoder.of());
    SetState<Integer> underlyingValue = underlying.state(namespace, valueTag);
    assertThat(underlyingValue.read(), emptyIterable());
    underlyingValue.add(1);
    assertThat(underlyingValue.read(), containsInAnyOrder(1));
    CopyOnAccessInMemoryStateInternals<String> internals = CopyOnAccessInMemoryStateInternals.withUnderlying(key, underlying);
    SetState<Integer> copyOnAccessState = internals.state(namespace, valueTag);
    assertThat(copyOnAccessState.read(), containsInAnyOrder(1));
    copyOnAccessState.add(4);
    assertThat(copyOnAccessState.read(), containsInAnyOrder(4, 1));
    assertThat(underlyingValue.read(), containsInAnyOrder(1));
    SetState<Integer> reReadUnderlyingValue = underlying.state(namespace, valueTag);
    assertThat(underlyingValue.read(), equalTo(reReadUnderlyingValue.read()));
}