Examples with StateNamespace org.apache.beam.runners.core.StateNamespace used on opensource projects

Example 31 with StateNamespace

use of org.apache.beam.runners.core.StateNamespace in project beam by apache.

the class CopyOnAccessInMemoryStateInternalsTest method testAccumulatorCombiningStateWithUnderlying.

@Test
public void testAccumulatorCombiningStateWithUnderlying() throws CannotProvideCoderException {
    CopyOnAccessInMemoryStateInternals<String> underlying = CopyOnAccessInMemoryStateInternals.withUnderlying(key, null);
    CombineFn<Long, long[], Long> sumLongFn = Sum.ofLongs();
    StateNamespace namespace = new StateNamespaceForTest("foo");
    CoderRegistry reg = pipeline.getCoderRegistry();
    StateTag<CombiningState<Long, long[], Long>> stateTag = StateTags.combiningValue("summer", sumLongFn.getAccumulatorCoder(reg, reg.getCoder(Long.class)), sumLongFn);
    GroupingState<Long, Long> underlyingValue = underlying.state(namespace, stateTag);
    assertThat(underlyingValue.read(), equalTo(0L));
    underlyingValue.add(1L);
    assertThat(underlyingValue.read(), equalTo(1L));
    CopyOnAccessInMemoryStateInternals<String> internals = CopyOnAccessInMemoryStateInternals.withUnderlying(key, underlying);
    GroupingState<Long, Long> copyOnAccessState = internals.state(namespace, stateTag);
    assertThat(copyOnAccessState.read(), equalTo(1L));
    copyOnAccessState.add(4L);
    assertThat(copyOnAccessState.read(), equalTo(5L));
    assertThat(underlyingValue.read(), equalTo(1L));
    GroupingState<Long, Long> reReadUnderlyingValue = underlying.state(namespace, stateTag);
    assertThat(underlyingValue.read(), equalTo(reReadUnderlyingValue.read()));
}

Example 32 with StateNamespace

use of org.apache.beam.runners.core.StateNamespace 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 33 with StateNamespace

use of org.apache.beam.runners.core.StateNamespace in project beam by apache.

the class CopyOnAccessInMemoryStateInternalsTest method testCommitWithClearedInUnderlying.

@Test
public void testCommitWithClearedInUnderlying() {
    CopyOnAccessInMemoryStateInternals<String> underlying = CopyOnAccessInMemoryStateInternals.withUnderlying(key, null);
    CopyOnAccessInMemoryStateInternals<String> secondUnderlying = spy(CopyOnAccessInMemoryStateInternals.withUnderlying(key, underlying));
    CopyOnAccessInMemoryStateInternals<String> internals = CopyOnAccessInMemoryStateInternals.withUnderlying(key, secondUnderlying);
    StateNamespace namespace = new StateNamespaceForTest("foo");
    StateTag<BagState<String>> bagTag = StateTags.bag("foo", StringUtf8Coder.of());
    BagState<String> stringBag = underlying.state(namespace, bagTag);
    assertThat(stringBag.read(), emptyIterable());
    stringBag.add("bar");
    stringBag.add("baz");
    stringBag.clear();
    // We should not read through the cleared bag
    secondUnderlying.commit();
    // Should not be visible
    stringBag.add("foo");
    internals.commit();
    BagState<String> internalsStringBag = internals.state(namespace, bagTag);
    assertThat(internalsStringBag.read(), emptyIterable());
    verify(secondUnderlying, never()).state(namespace, bagTag);
    assertThat(internals.isEmpty(), is(false));
}

Example 34 with StateNamespace

use of org.apache.beam.runners.core.StateNamespace in project beam by apache.

the class CopyOnAccessInMemoryStateInternalsTest method testMapStateWithUnderlying.

@Test
public void testMapStateWithUnderlying() {
    CopyOnAccessInMemoryStateInternals<String> underlying = CopyOnAccessInMemoryStateInternals.withUnderlying(key, null);
    StateNamespace namespace = new StateNamespaceForTest("foo");
    StateTag<MapState<String, Integer>> valueTag = StateTags.map("foo", StringUtf8Coder.of(), VarIntCoder.of());
    MapState<String, Integer> underlyingValue = underlying.state(namespace, valueTag);
    assertThat(underlyingValue.entries().read(), emptyIterable());
    underlyingValue.put("hello", 1);
    assertThat(underlyingValue.get("hello").read(), equalTo(1));
    CopyOnAccessInMemoryStateInternals<String> internals = CopyOnAccessInMemoryStateInternals.withUnderlying(key, underlying);
    MapState<String, Integer> copyOnAccessState = internals.state(namespace, valueTag);
    assertThat(copyOnAccessState.get("hello").read(), equalTo(1));
    copyOnAccessState.put("world", 4);
    assertThat(copyOnAccessState.get("hello").read(), equalTo(1));
    assertThat(copyOnAccessState.get("world").read(), equalTo(4));
    assertThat(underlyingValue.get("hello").read(), equalTo(1));
    assertNull(underlyingValue.get("world").read());
    MapState<String, Integer> reReadUnderlyingValue = underlying.state(namespace, valueTag);
    assertThat(underlyingValue.entries().read(), equalTo(reReadUnderlyingValue.entries().read()));
}

Example 35 with StateNamespace

use of org.apache.beam.runners.core.StateNamespace in project beam by apache.

the class FlinkStatefulDoFnFunction method fireTimer.

private void fireTimer(final K key, TimerInternals.TimerData timer, DoFnRunner<KV<K, V>, OutputT> doFnRunner) {
    StateNamespace namespace = timer.getNamespace();
    checkArgument(namespace instanceof StateNamespaces.WindowNamespace);
    BoundedWindow window = ((StateNamespaces.WindowNamespace) namespace).getWindow();
    doFnRunner.onTimer(timer.getTimerId(), timer.getTimerFamilyId(), key, window, timer.getTimestamp(), timer.getOutputTimestamp(), timer.getDomain());
}