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

Example 1 with DefaultKeyedStateStore

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

the class KeyedStateInputFormat method open.

@Override
@SuppressWarnings("unchecked")
public void open(KeyGroupRangeInputSplit split) throws IOException {
    registry = new CloseableRegistry();
    final StreamOperatorStateContext context = new StreamOperatorContextBuilder(getRuntimeContext(), configuration, operatorState, split, registry, stateBackend).withMaxParallelism(split.getNumKeyGroups()).withKey(operator, operator.getKeyType().createSerializer(getRuntimeContext().getExecutionConfig())).build(LOG);
    AbstractKeyedStateBackend<K> keyedStateBackend = (AbstractKeyedStateBackend<K>) context.keyedStateBackend();
    final DefaultKeyedStateStore keyedStateStore = new DefaultKeyedStateStore(keyedStateBackend, getRuntimeContext().getExecutionConfig());
    SavepointRuntimeContext ctx = new SavepointRuntimeContext(getRuntimeContext(), keyedStateStore);
    InternalTimeServiceManager<K> timeServiceManager = (InternalTimeServiceManager<K>) context.internalTimerServiceManager();
    try {
        operator.setup(getRuntimeContext().getExecutionConfig(), keyedStateBackend, timeServiceManager, ctx);
        operator.open();
        keysAndNamespaces = operator.getKeysAndNamespaces(ctx);
    } catch (Exception e) {
        throw new IOException("Failed to restore timer state", e);
    }
}

Example 2 with DefaultKeyedStateStore

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

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

the class AbstractStreamOperator method initKeyedState.

private void initKeyedState() {
    try {
        TypeSerializer<Object> keySerializer = config.getStateKeySerializer(getUserCodeClassloader());
        // create a keyed state backend if there is keyed state, as indicated by the presence of a key serializer
        if (null != keySerializer) {
            KeyGroupRange subTaskKeyGroupRange = KeyGroupRangeAssignment.computeKeyGroupRangeForOperatorIndex(container.getEnvironment().getTaskInfo().getMaxNumberOfParallelSubtasks(), container.getEnvironment().getTaskInfo().getNumberOfParallelSubtasks(), container.getEnvironment().getTaskInfo().getIndexOfThisSubtask());
            this.keyedStateBackend = container.createKeyedStateBackend(keySerializer, // The maximum parallelism == number of key group
            container.getEnvironment().getTaskInfo().getMaxNumberOfParallelSubtasks(), subTaskKeyGroupRange);
            this.keyedStateStore = new DefaultKeyedStateStore(keyedStateBackend, getExecutionConfig());
        }
    } catch (Exception e) {
        throw new IllegalStateException("Could not initialize keyed state backend.", e);
    }
}

Example 4 with DefaultKeyedStateStore

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

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

the class StreamingRuntimeContextTest method createMapPlainMockOp.

@SuppressWarnings("unchecked")
private static AbstractStreamOperator<?> createMapPlainMockOp() 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<MapState<Integer, String>>() {

        @Override
        public MapState<Integer, String> answer(InvocationOnMock invocationOnMock) throws Throwable {
            MapStateDescriptor<Integer, String> descr = (MapStateDescriptor<Integer, 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(MapStateDescriptor.class));
    when(operatorMock.getKeyedStateStore()).thenReturn(keyedStateStore);
    when(operatorMock.getOperatorID()).thenReturn(new OperatorID());
    when(operatorMock.getProcessingTimeService()).thenReturn(new TestProcessingTimeService());
    return operatorMock;
}