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

Example 11 with BagState

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

the class FlinkKeyGroupStateInternalsTest method testKeyGroupAndCheckpoint.

@Test
public void testKeyGroupAndCheckpoint() throws Exception {
    // assign to keyGroup 0
    ByteBuffer key0 = ByteBuffer.wrap(CoderUtils.encodeToByteArray(StringUtf8Coder.of(), "11111111"));
    // assign to keyGroup 1
    ByteBuffer key1 = ByteBuffer.wrap(CoderUtils.encodeToByteArray(StringUtf8Coder.of(), "22222222"));
    FlinkKeyGroupStateInternals<String> allState;
    {
        KeyedStateBackend keyedStateBackend = getKeyedStateBackend(2, new KeyGroupRange(0, 1));
        allState = new FlinkKeyGroupStateInternals<>(StringUtf8Coder.of(), keyedStateBackend);
        BagState<String> valueForNamespace0 = allState.state(NAMESPACE_1, STRING_BAG_ADDR);
        BagState<String> valueForNamespace1 = allState.state(NAMESPACE_2, STRING_BAG_ADDR);
        keyedStateBackend.setCurrentKey(key0);
        valueForNamespace0.add("0");
        valueForNamespace1.add("2");
        keyedStateBackend.setCurrentKey(key1);
        valueForNamespace0.add("1");
        valueForNamespace1.add("3");
        assertThat(valueForNamespace0.read(), Matchers.containsInAnyOrder("0", "1"));
        assertThat(valueForNamespace1.read(), Matchers.containsInAnyOrder("2", "3"));
    }
    ClassLoader classLoader = FlinkKeyGroupStateInternalsTest.class.getClassLoader();
    // 1. scale up
    ByteArrayOutputStream out0 = new ByteArrayOutputStream();
    allState.snapshotKeyGroupState(0, new DataOutputStream(out0));
    DataInputStream in0 = new DataInputStream(new ByteArrayInputStream(out0.toByteArray()));
    {
        KeyedStateBackend keyedStateBackend = getKeyedStateBackend(2, new KeyGroupRange(0, 0));
        FlinkKeyGroupStateInternals<String> state0 = new FlinkKeyGroupStateInternals<>(StringUtf8Coder.of(), keyedStateBackend);
        state0.restoreKeyGroupState(0, in0, classLoader);
        BagState<String> valueForNamespace0 = state0.state(NAMESPACE_1, STRING_BAG_ADDR);
        BagState<String> valueForNamespace1 = state0.state(NAMESPACE_2, STRING_BAG_ADDR);
        assertThat(valueForNamespace0.read(), Matchers.containsInAnyOrder("0"));
        assertThat(valueForNamespace1.read(), Matchers.containsInAnyOrder("2"));
    }
    ByteArrayOutputStream out1 = new ByteArrayOutputStream();
    allState.snapshotKeyGroupState(1, new DataOutputStream(out1));
    DataInputStream in1 = new DataInputStream(new ByteArrayInputStream(out1.toByteArray()));
    {
        KeyedStateBackend keyedStateBackend = getKeyedStateBackend(2, new KeyGroupRange(1, 1));
        FlinkKeyGroupStateInternals<String> state1 = new FlinkKeyGroupStateInternals<>(StringUtf8Coder.of(), keyedStateBackend);
        state1.restoreKeyGroupState(1, in1, classLoader);
        BagState<String> valueForNamespace0 = state1.state(NAMESPACE_1, STRING_BAG_ADDR);
        BagState<String> valueForNamespace1 = state1.state(NAMESPACE_2, STRING_BAG_ADDR);
        assertThat(valueForNamespace0.read(), Matchers.containsInAnyOrder("1"));
        assertThat(valueForNamespace1.read(), Matchers.containsInAnyOrder("3"));
    }
    // 2. scale down
    {
        KeyedStateBackend keyedStateBackend = getKeyedStateBackend(2, new KeyGroupRange(0, 1));
        FlinkKeyGroupStateInternals<String> newAllState = new FlinkKeyGroupStateInternals<>(StringUtf8Coder.of(), keyedStateBackend);
        in0.reset();
        in1.reset();
        newAllState.restoreKeyGroupState(0, in0, classLoader);
        newAllState.restoreKeyGroupState(1, in1, classLoader);
        BagState<String> valueForNamespace0 = newAllState.state(NAMESPACE_1, STRING_BAG_ADDR);
        BagState<String> valueForNamespace1 = newAllState.state(NAMESPACE_2, STRING_BAG_ADDR);
        assertThat(valueForNamespace0.read(), Matchers.containsInAnyOrder("0", "1"));
        assertThat(valueForNamespace1.read(), Matchers.containsInAnyOrder("2", "3"));
    }
}

Example 12 with BagState

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

the class EvaluationContextTest method getExecutionContextSameStepSameKeyState.

@Test
public void getExecutionContextSameStepSameKeyState() {
    DirectExecutionContext fooContext = context.getExecutionContext(createdProducer, StructuralKey.of("foo", StringUtf8Coder.of()));
    StateTag<BagState<Integer>> intBag = StateTags.bag("myBag", VarIntCoder.of());
    DirectStepContext stepContext = fooContext.getStepContext("s1");
    stepContext.stateInternals().state(StateNamespaces.global(), intBag).add(1);
    context.handleResult(ImmutableListBundleFactory.create().createKeyedBundle(StructuralKey.of("foo", StringUtf8Coder.of()), created).commit(Instant.now()), ImmutableList.<TimerData>of(), StepTransformResult.withoutHold(createdProducer).withState(stepContext.commitState()).build());
    DirectExecutionContext secondFooContext = context.getExecutionContext(createdProducer, StructuralKey.of("foo", StringUtf8Coder.of()));
    assertThat(secondFooContext.getStepContext("s1").stateInternals().state(StateNamespaces.global(), intBag).read(), contains(1));
}

Example 13 with BagState

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

the class CopyOnAccessInMemoryStateInternalsTest method testCommitWithOnlyClearedValuesIsEmpty.

@Test
public void testCommitWithOnlyClearedValuesIsEmpty() {
    CopyOnAccessInMemoryStateInternals<String> internals = CopyOnAccessInMemoryStateInternals.withUnderlying(key, null);
    StateNamespace namespace = new StateNamespaceForTest("foo");
    StateTag<BagState<String>> bagTag = StateTags.bag("foo", StringUtf8Coder.of());
    BagState<String> stringBag = internals.state(namespace, bagTag);
    assertThat(stringBag.read(), emptyIterable());
    stringBag.add("foo");
    stringBag.clear();
    internals.commit();
    assertThat(internals.isEmpty(), is(true));
}

Example 14 with BagState

use of org.apache.beam.sdk.state.BagState 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 15 with BagState

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

the class CopyOnAccessInMemoryStateInternalsTest method testCommitWithAddedUnderlying.

@Test
public void testCommitWithAddedUnderlying() {
    CopyOnAccessInMemoryStateInternals<String> underlying = CopyOnAccessInMemoryStateInternals.withUnderlying(key, null);
    CopyOnAccessInMemoryStateInternals<String> internals = CopyOnAccessInMemoryStateInternals.withUnderlying(key, underlying);
    internals.commit();
    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");
    BagState<String> internalState = internals.state(namespace, bagTag);
    assertThat(internalState.read(), emptyIterable());
    BagState<String> reReadUnderlyingState = underlying.state(namespace, bagTag);
    assertThat(reReadUnderlyingState.read(), containsInAnyOrder("bar", "baz"));
}