Examples with FlinkKeyGroupStateInternals org.apache.beam.runners.flink.translation.wrappers.streaming.state.FlinkKeyGroupStateInternals used on opensource projects

Example 1 with FlinkKeyGroupStateInternals

use of org.apache.beam.runners.flink.translation.wrappers.streaming.state.FlinkKeyGroupStateInternals 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"));
    }
}