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

Example 16 with SamzaPipelineOptions

use of org.apache.beam.runners.samza.SamzaPipelineOptions in project beam by apache.

the class SamzaTimerInternalsFactoryTest method testEventTimeTimers.

@Test
public void testEventTimeTimers() {
    final SamzaPipelineOptions pipelineOptions = PipelineOptionsFactory.create().as(SamzaPipelineOptions.class);
    final KeyValueStore<ByteArray, StateValue<?>> store = createStore();
    final SamzaTimerInternalsFactory<String> timerInternalsFactory = createTimerInternalsFactory(null, "timer", pipelineOptions, store);
    final StateNamespace nameSpace = StateNamespaces.global();
    final TimerInternals timerInternals = timerInternalsFactory.timerInternalsForKey("testKey");
    final TimerInternals.TimerData timer1 = TimerInternals.TimerData.of("timer1", nameSpace, new Instant(10), new Instant(10), TimeDomain.EVENT_TIME);
    timerInternals.setTimer(timer1);
    final TimerInternals.TimerData timer2 = TimerInternals.TimerData.of("timer2", nameSpace, new Instant(100), new Instant(100), TimeDomain.EVENT_TIME);
    timerInternals.setTimer(timer2);
    timerInternalsFactory.setInputWatermark(new Instant(5));
    Collection<KeyedTimerData<String>> readyTimers = timerInternalsFactory.removeReadyTimers();
    assertTrue(readyTimers.isEmpty());
    timerInternalsFactory.setInputWatermark(new Instant(20));
    readyTimers = timerInternalsFactory.removeReadyTimers();
    assertEquals(1, readyTimers.size());
    assertEquals(timer1, readyTimers.iterator().next().getTimerData());
    timerInternalsFactory.setInputWatermark(new Instant(150));
    readyTimers = timerInternalsFactory.removeReadyTimers();
    assertEquals(1, readyTimers.size());
    assertEquals(timer2, readyTimers.iterator().next().getTimerData());
    store.close();
}

Example 17 with SamzaPipelineOptions

use of org.apache.beam.runners.samza.SamzaPipelineOptions in project beam by apache.

the class SamzaTimerInternalsFactoryTest method testOverride.

@Test
public void testOverride() {
    final SamzaPipelineOptions pipelineOptions = PipelineOptionsFactory.create().as(SamzaPipelineOptions.class);
    KeyValueStore<ByteArray, StateValue<?>> store = createStore();
    final SamzaTimerInternalsFactory<String> timerInternalsFactory = createTimerInternalsFactory(null, "timer", pipelineOptions, store);
    final StateNamespace nameSpace = StateNamespaces.global();
    final TimerInternals timerInternals = timerInternalsFactory.timerInternalsForKey("testKey");
    final TimerInternals.TimerData timer1 = TimerInternals.TimerData.of("timerId", nameSpace, new Instant(10), new Instant(10), TimeDomain.EVENT_TIME);
    timerInternals.setTimer(timer1);
    // this timer should override the first timer
    final TimerInternals.TimerData timer2 = TimerInternals.TimerData.of("timerId", nameSpace, new Instant(100), new Instant(100), TimeDomain.EVENT_TIME);
    timerInternals.setTimer(timer2);
    final TimerInternals.TimerData timer3 = TimerInternals.TimerData.of("timerId2", nameSpace, new Instant(200), new Instant(200), TimeDomain.EVENT_TIME);
    timerInternals.setTimer(timer3);
    // this timer shouldn't override since it has a different id
    timerInternalsFactory.setInputWatermark(new Instant(50));
    Collection<KeyedTimerData<String>> readyTimers = timerInternalsFactory.removeReadyTimers();
    assertEquals(0, readyTimers.size());
    timerInternalsFactory.setInputWatermark(new Instant(150));
    readyTimers = timerInternalsFactory.removeReadyTimers();
    assertEquals(1, readyTimers.size());
    timerInternalsFactory.setInputWatermark(new Instant(250));
    readyTimers = timerInternalsFactory.removeReadyTimers();
    assertEquals(1, readyTimers.size());
    store.close();
}

Example 18 with SamzaPipelineOptions

use of org.apache.beam.runners.samza.SamzaPipelineOptions 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));
}

Example 19 with SamzaPipelineOptions

use of org.apache.beam.runners.samza.SamzaPipelineOptions in project beam by apache.

the class ConfigGeneratorTest method testUserStoreConfig.

@Test
public void testUserStoreConfig() {
    SamzaPipelineOptions options = PipelineOptionsFactory.create().as(SamzaPipelineOptions.class);
    options.setJobName("TestStoreConfig");
    options.setRunner(SamzaRunner.class);
    Pipeline pipeline = Pipeline.create(options);
    pipeline.apply(Create.empty(TypeDescriptors.kvs(TypeDescriptors.strings(), TypeDescriptors.strings()))).apply(ParDo.of(new DoFn<KV<String, String>, Void>() {

        private static final String testState = "testState";

        @StateId(testState)
        private final StateSpec<ValueState<Integer>> state = StateSpecs.value();

        @ProcessElement
        public void processElement(ProcessContext context, @StateId(testState) ValueState<Integer> state) {
        }
    }));
    final Map<PValue, String> idMap = PViewToIdMapper.buildIdMap(pipeline);
    final ConfigBuilder configBuilder = new ConfigBuilder(options);
    SamzaPipelineTranslator.createConfig(pipeline, options, idMap, configBuilder);
    final Config config = configBuilder.build();
    assertEquals(RocksDbKeyValueStorageEngineFactory.class.getName(), config.get("stores.testState.factory"));
    assertEquals("byteArraySerde", config.get("stores.testState.key.serde"));
    assertEquals("stateValueSerde", config.get("stores.testState.msg.serde"));
    assertNull(config.get("stores.testState.changelog"));
    options.setStateDurable(true);
    SamzaPipelineTranslator.createConfig(pipeline, options, idMap, configBuilder);
    final Config config2 = configBuilder.build();
    assertEquals("TestStoreConfig-1-testState-changelog", config2.get("stores.testState.changelog"));
}

Example 20 with SamzaPipelineOptions

use of org.apache.beam.runners.samza.SamzaPipelineOptions in project beam by apache.

the class ConfigGeneratorTest method testStatelessBeamStoreConfig.

@Test
public void testStatelessBeamStoreConfig() {
    SamzaPipelineOptions options = PipelineOptionsFactory.create().as(SamzaPipelineOptions.class);
    options.setJobName("TestStoreConfig");
    options.setRunner(SamzaRunner.class);
    Pipeline pipeline = Pipeline.create(options);
    pipeline.apply(Impulse.create()).apply(Filter.by(Objects::nonNull));
    pipeline.replaceAll(SamzaTransformOverrides.getDefaultOverrides());
    final Map<PValue, String> idMap = PViewToIdMapper.buildIdMap(pipeline);
    final ConfigBuilder configBuilder = new ConfigBuilder(options);
    SamzaPipelineTranslator.createConfig(pipeline, options, idMap, configBuilder);
    final Config config = configBuilder.build();
    assertEquals(InMemoryKeyValueStorageEngineFactory.class.getName(), config.get("stores.beamStore.factory"));
    assertEquals("byteArraySerde", config.get("stores.beamStore.key.serde"));
    assertEquals("stateValueSerde", config.get("stores.beamStore.msg.serde"));
    assertNull(config.get("stores.beamStore.changelog"));
    options.setStateDurable(true);
    SamzaPipelineTranslator.createConfig(pipeline, options, idMap, configBuilder);
    final Config config2 = configBuilder.build();
    // For stateless jobs, ignore state durable pipeline option.
    assertNull(config2.get("stores.beamStore.changelog"));
}