Examples with SamzaSetState org.apache.beam.runners.samza.state.SamzaSetState used on opensource projects

Example 1 with SamzaSetState

use of org.apache.beam.runners.samza.state.SamzaSetState in project beam by apache.

the class SamzaStoreStateInternalsTest method testIteratorClosed.

@Test
public void testIteratorClosed() {
    final String stateId = "foo";
    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());

        @ProcessElement
        public void processElement(ProcessContext c, @StateId(stateId) SetState<Integer> setState) {
            SamzaSetState<Integer> state = (SamzaSetState<Integer>) setState;
            state.add(c.element().getValue());
            // the iterator for size needs to be closed
            int size = Iterators.size(state.readIterator().read());
            if (size > 1) {
                final Iterator<Integer> iterator = state.readIterator().read();
                assertTrue(iterator.hasNext());
                // this iterator should be closed too
                iterator.next();
            }
        }
    };
    pipeline.apply(Create.of(KV.of("hello", 97), KV.of("hello", 42), KV.of("hello", 42), KV.of("hello", 12))).apply(ParDo.of(fn));
    SamzaPipelineOptions options = PipelineOptionsFactory.create().as(SamzaPipelineOptions.class);
    options.setRunner(TestSamzaRunner.class);
    Map<String, String> configs = new HashMap<>(ConfigBuilder.localRunConfig());
    configs.put("stores.foo.factory", TestStorageEngine.class.getName());
    pipeline.getOptions().as(SamzaPipelineOptions.class).setConfigOverride(configs);
    pipeline.run();
    // The test code creates 7 underlying iterators, and 1 more is created during state.clear()
    // Verify all of them are closed
    assertEquals(8, TestStore.iterators.size());
    TestStore.iterators.forEach(iter -> assertTrue(iter.closed));
}