Examples with KeyedStateBackend org.apache.flink.runtime.state.KeyedStateBackend used on opensource projects

Example 1 with KeyedStateBackend

use of org.apache.flink.runtime.state.KeyedStateBackend in project flink by apache.

the class StreamingRuntimeContextTest method createListPlainMockOp.

@SuppressWarnings("unchecked")
private static AbstractStreamOperator<?> createListPlainMockOp() throws Exception {
    AbstractStreamOperator<?> operatorMock = mock(AbstractStreamOperator.class);
    ExecutionConfig config = new ExecutionConfig();
    KeyedStateBackend keyedStateBackend = mock(KeyedStateBackend.class);
    DefaultKeyedStateStore keyedStateStore = new DefaultKeyedStateStore(keyedStateBackend, config);
    when(operatorMock.getExecutionConfig()).thenReturn(config);
    doAnswer(new Answer<ListState<String>>() {

        @Override
        public ListState<String> answer(InvocationOnMock invocationOnMock) throws Throwable {
            ListStateDescriptor<String> descr = (ListStateDescriptor<String>) invocationOnMock.getArguments()[2];
            AbstractKeyedStateBackend<Integer> backend = new MemoryStateBackend().createKeyedStateBackend(new DummyEnvironment("test_task", 1, 0), new JobID(), "test_op", IntSerializer.INSTANCE, 1, new KeyGroupRange(0, 0), new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID()), TtlTimeProvider.DEFAULT, new UnregisteredMetricsGroup(), Collections.emptyList(), new CloseableRegistry());
            backend.setCurrentKey(0);
            return backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, descr);
        }
    }).when(keyedStateBackend).getPartitionedState(Matchers.any(), any(TypeSerializer.class), any(ListStateDescriptor.class));
    when(operatorMock.getKeyedStateStore()).thenReturn(keyedStateStore);
    when(operatorMock.getOperatorID()).thenReturn(new OperatorID());
    return operatorMock;
}

Example 2 with KeyedStateBackend

use of org.apache.flink.runtime.state.KeyedStateBackend in project beam by apache.

the class DoFnOperator method initializeState.

@Override
public void initializeState(StateInitializationContext context) throws Exception {
    super.initializeState(context);
    ListStateDescriptor<WindowedValue<InputT>> pushedBackStateDescriptor = new ListStateDescriptor<>("pushed-back-elements", new CoderTypeSerializer<>(windowedInputCoder, serializedOptions));
    if (keySelector != null) {
        pushedBackElementsHandler = KeyedPushedBackElementsHandler.create(keySelector, getKeyedStateBackend(), pushedBackStateDescriptor);
    } else {
        ListState<WindowedValue<InputT>> listState = getOperatorStateBackend().getListState(pushedBackStateDescriptor);
        pushedBackElementsHandler = NonKeyedPushedBackElementsHandler.create(listState);
    }
    currentInputWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE.getMillis();
    currentSideInputWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE.getMillis();
    currentOutputWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE.getMillis();
    sideInputReader = NullSideInputReader.of(sideInputs);
    if (!sideInputs.isEmpty()) {
        FlinkBroadcastStateInternals sideInputStateInternals = new FlinkBroadcastStateInternals<>(getContainingTask().getIndexInSubtaskGroup(), getOperatorStateBackend(), serializedOptions);
        sideInputHandler = new SideInputHandler(sideInputs, sideInputStateInternals);
        sideInputReader = sideInputHandler;
        Stream<WindowedValue<InputT>> pushedBack = pushedBackElementsHandler.getElements();
        long min = pushedBack.map(v -> v.getTimestamp().getMillis()).reduce(Long.MAX_VALUE, Math::min);
        pushedBackWatermark = min;
    } else {
        pushedBackWatermark = Long.MAX_VALUE;
    }
    // StatefulPardo or WindowDoFn
    if (keyCoder != null) {
        keyedStateInternals = new FlinkStateInternals<>((KeyedStateBackend) getKeyedStateBackend(), keyCoder, serializedOptions);
        if (timerService == null) {
            timerService = getInternalTimerService("beam-timer", new CoderTypeSerializer<>(timerCoder, serializedOptions), this);
        }
        timerInternals = new FlinkTimerInternals();
        timeServiceManagerCompat = getTimeServiceManagerCompat();
    }
    outputManager = outputManagerFactory.create(output, getLockToAcquireForStateAccessDuringBundles(), getOperatorStateBackend());
}

Example 3 with KeyedStateBackend

use of org.apache.flink.runtime.state.KeyedStateBackend 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 4 with KeyedStateBackend

use of org.apache.flink.runtime.state.KeyedStateBackend in project flink by apache.

the class StreamingRuntimeContextTest method createDescriptorCapturingMockOp.

// ------------------------------------------------------------------------
//  
// ------------------------------------------------------------------------
@SuppressWarnings("unchecked")
private static AbstractStreamOperator<?> createDescriptorCapturingMockOp(final AtomicReference<Object> ref, final ExecutionConfig config) throws Exception {
    AbstractStreamOperator<?> operatorMock = mock(AbstractStreamOperator.class);
    KeyedStateBackend keyedStateBackend = mock(KeyedStateBackend.class);
    DefaultKeyedStateStore keyedStateStore = new DefaultKeyedStateStore(keyedStateBackend, config);
    when(operatorMock.getExecutionConfig()).thenReturn(config);
    doAnswer(new Answer<Object>() {

        @Override
        public Object answer(InvocationOnMock invocationOnMock) throws Throwable {
            ref.set(invocationOnMock.getArguments()[2]);
            return null;
        }
    }).when(keyedStateBackend).getPartitionedState(Matchers.any(), any(TypeSerializer.class), any(StateDescriptor.class));
    when(operatorMock.getKeyedStateStore()).thenReturn(keyedStateStore);
    return operatorMock;
}

Example 5 with KeyedStateBackend

use of org.apache.flink.runtime.state.KeyedStateBackend in project beam by apache.

the class ExecutableStageDoFnOperatorTest method testEnsureStateCleanupWithKeyedInputCleanupTimer.

@SuppressWarnings("LockNotBeforeTry")
@Test
public void testEnsureStateCleanupWithKeyedInputCleanupTimer() {
    InMemoryTimerInternals inMemoryTimerInternals = new InMemoryTimerInternals();
    KeyedStateBackend keyedStateBackend = Mockito.mock(KeyedStateBackend.class);
    Lock stateBackendLock = Mockito.mock(Lock.class);
    StringUtf8Coder keyCoder = StringUtf8Coder.of();
    IntervalWindow window = new IntervalWindow(new Instant(0), new Instant(10));
    Coder<IntervalWindow> windowCoder = IntervalWindow.getCoder();
    // Test that cleanup timer is set correctly
    ExecutableStageDoFnOperator.CleanupTimer cleanupTimer = new ExecutableStageDoFnOperator.CleanupTimer<>(inMemoryTimerInternals, stateBackendLock, WindowingStrategy.globalDefault(), keyCoder, windowCoder, keyedStateBackend);
    cleanupTimer.setForWindow(KV.of("key", "string"), window);
    Mockito.verify(stateBackendLock).lock();
    ByteBuffer key = FlinkKeyUtils.encodeKey("key", keyCoder);
    Mockito.verify(keyedStateBackend).setCurrentKey(key);
    assertThat(inMemoryTimerInternals.getNextTimer(TimeDomain.EVENT_TIME), is(window.maxTimestamp().plus(Duration.millis(1))));
    Mockito.verify(stateBackendLock).unlock();
}