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

Example 16 with BagState

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

the class ParDoTest method testBagState.

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

        @StateId(stateId)
        private final StateSpec<BagState<Integer>> bufferState = StateSpecs.bag(VarIntCoder.of());

        @ProcessElement
        public void processElement(ProcessContext c, @StateId(stateId) BagState<Integer> state) {
            Iterable<Integer> currentValue = state.read();
            state.add(c.element().getValue());
            if (Iterables.size(state.read()) >= 4) {
                List<Integer> sorted = Lists.newArrayList(currentValue);
                Collections.sort(sorted);
                c.output(sorted);
            }
        }
    };
    PCollection<List<Integer>> output = pipeline.apply(Create.of(KV.of("hello", 97), KV.of("hello", 42), KV.of("hello", 84), KV.of("hello", 12))).apply(ParDo.of(fn));
    PAssert.that(output).containsInAnyOrder(Lists.newArrayList(12, 42, 84, 97));
    pipeline.run();
}