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

Example 36 with KeyGroupRange

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

the class StateInitializationContextImplTest method setUp.

@Before
public void setUp() throws Exception {
    this.writtenKeyGroups = 0;
    this.writtenOperatorStates = new HashSet<>();
    this.closableRegistry = new CloseableRegistry();
    ByteArrayOutputStreamWithPos out = new ByteArrayOutputStreamWithPos(64);
    List<KeyedStateHandle> keyedStateHandles = new ArrayList<>(NUM_HANDLES);
    int prev = 0;
    for (int i = 0; i < NUM_HANDLES; ++i) {
        out.reset();
        int size = i % 4;
        int end = prev + size;
        DataOutputView dov = new DataOutputViewStreamWrapper(out);
        KeyGroupRangeOffsets offsets = new KeyGroupRangeOffsets(i == 9 ? KeyGroupRange.EMPTY_KEY_GROUP_RANGE : new KeyGroupRange(prev, end));
        prev = end + 1;
        for (int kg : offsets.getKeyGroupRange()) {
            offsets.setKeyGroupOffset(kg, out.getPosition());
            dov.writeInt(kg);
            ++writtenKeyGroups;
        }
        KeyedStateHandle handle = new KeyGroupsStateHandle(offsets, new ByteStateHandleCloseChecking("kg-" + i, out.toByteArray()));
        keyedStateHandles.add(handle);
    }
    List<OperatorStateHandle> operatorStateHandles = new ArrayList<>(NUM_HANDLES);
    for (int i = 0; i < NUM_HANDLES; ++i) {
        int size = i % 4;
        out.reset();
        DataOutputView dov = new DataOutputViewStreamWrapper(out);
        LongArrayList offsets = new LongArrayList(size);
        for (int s = 0; s < size; ++s) {
            offsets.add(out.getPosition());
            int val = i * NUM_HANDLES + s;
            dov.writeInt(val);
            writtenOperatorStates.add(val);
        }
        Map<String, OperatorStateHandle.StateMetaInfo> offsetsMap = new HashMap<>();
        offsetsMap.put(DefaultOperatorStateBackend.DEFAULT_OPERATOR_STATE_NAME, new OperatorStateHandle.StateMetaInfo(offsets.toArray(), OperatorStateHandle.Mode.SPLIT_DISTRIBUTE));
        OperatorStateHandle operatorStateHandle = new OperatorStreamStateHandle(offsetsMap, new ByteStateHandleCloseChecking("os-" + i, out.toByteArray()));
        operatorStateHandles.add(operatorStateHandle);
    }
    OperatorSubtaskState operatorSubtaskState = OperatorSubtaskState.builder().setRawOperatorState(new StateObjectCollection<>(operatorStateHandles)).setRawKeyedState(new StateObjectCollection<>(keyedStateHandles)).build();
    OperatorID operatorID = new OperatorID();
    TaskStateSnapshot taskStateSnapshot = new TaskStateSnapshot();
    taskStateSnapshot.putSubtaskStateByOperatorID(operatorID, operatorSubtaskState);
    JobManagerTaskRestore jobManagerTaskRestore = new JobManagerTaskRestore(0L, taskStateSnapshot);
    TaskStateManager manager = new TaskStateManagerImpl(new JobID(), new ExecutionAttemptID(), new TestTaskLocalStateStore(), new InMemoryStateChangelogStorage(), jobManagerTaskRestore, mock(CheckpointResponder.class));
    DummyEnvironment environment = new DummyEnvironment("test", 1, 0, prev);
    environment.setTaskStateManager(manager);
    StateBackend stateBackend = new MemoryStateBackend(1024);
    StreamTaskStateInitializer streamTaskStateManager = new StreamTaskStateInitializerImpl(environment, stateBackend, TtlTimeProvider.DEFAULT, new InternalTimeServiceManager.Provider() {

        @Override
        public <K> InternalTimeServiceManager<K> create(CheckpointableKeyedStateBackend<K> keyedStatedBackend, ClassLoader userClassloader, KeyContext keyContext, ProcessingTimeService processingTimeService, Iterable<KeyGroupStatePartitionStreamProvider> rawKeyedStates) throws Exception {
            // stream.
            return null;
        }
    });
    AbstractStreamOperator<?> mockOperator = mock(AbstractStreamOperator.class);
    when(mockOperator.getOperatorID()).thenReturn(operatorID);
    StreamOperatorStateContext stateContext = streamTaskStateManager.streamOperatorStateContext(operatorID, "TestOperatorClass", mock(ProcessingTimeService.class), mockOperator, // consumed by the timer service.
    IntSerializer.INSTANCE, closableRegistry, new UnregisteredMetricsGroup(), 1.0, false);
    OptionalLong restoredCheckpointId = stateContext.getRestoredCheckpointId();
    this.initializationContext = new StateInitializationContextImpl(restoredCheckpointId.isPresent() ? restoredCheckpointId.getAsLong() : null, stateContext.operatorStateBackend(), mock(KeyedStateStore.class), stateContext.rawKeyedStateInputs(), stateContext.rawOperatorStateInputs());
}

Example 37 with KeyGroupRange

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

the class StateSnapshotContextSynchronousImplTest method testStreamClosingWhenClosing.

/**
 * Tests that closing the StateSnapshotContextSynchronousImpl will also close the associated
 * output streams.
 */
@Test
public void testStreamClosingWhenClosing() throws Exception {
    long checkpointId = 42L;
    long checkpointTimestamp = 1L;
    CheckpointStateOutputStream outputStream1 = mock(CheckpointStateOutputStream.class);
    CheckpointStateOutputStream outputStream2 = mock(CheckpointStateOutputStream.class);
    CheckpointStreamFactory streamFactory = mock(CheckpointStreamFactory.class);
    when(streamFactory.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE)).thenReturn(outputStream1, outputStream2);
    InsightCloseableRegistry closableRegistry = new InsightCloseableRegistry();
    KeyGroupRange keyGroupRange = new KeyGroupRange(0, 2);
    StateSnapshotContextSynchronousImpl context = new StateSnapshotContextSynchronousImpl(checkpointId, checkpointTimestamp, streamFactory, keyGroupRange, closableRegistry);
    // creating the output streams
    context.getRawKeyedOperatorStateOutput();
    context.getRawOperatorStateOutput();
    verify(streamFactory, times(2)).createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);
    assertEquals(2, closableRegistry.size());
    assertTrue(closableRegistry.contains(outputStream1));
    assertTrue(closableRegistry.contains(outputStream2));
    context.getKeyedStateStreamFuture().run();
    context.getOperatorStateStreamFuture().run();
    verify(outputStream1).closeAndGetHandle();
    verify(outputStream2).closeAndGetHandle();
    assertEquals(0, closableRegistry.size());
}

Example 38 with KeyGroupRange

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

the class StateSnapshotContextSynchronousImplTest method setUp.

@Before
public void setUp() throws Exception {
    CloseableRegistry closableRegistry = new CloseableRegistry();
    CheckpointStreamFactory streamFactory = new MemCheckpointStreamFactory(1024);
    KeyGroupRange keyGroupRange = new KeyGroupRange(0, 2);
    this.snapshotContext = new StateSnapshotContextSynchronousImpl(42, 4711, streamFactory, keyGroupRange, closableRegistry);
}

Example 39 with KeyGroupRange

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

the class StateSnapshotContextSynchronousImplTest method testStreamClosingExceptionally.

@Test
public void testStreamClosingExceptionally() throws Exception {
    long checkpointId = 42L;
    long checkpointTimestamp = 1L;
    CheckpointStateOutputStream outputStream1 = mock(CheckpointStateOutputStream.class);
    CheckpointStateOutputStream outputStream2 = mock(CheckpointStateOutputStream.class);
    CheckpointStreamFactory streamFactory = mock(CheckpointStreamFactory.class);
    when(streamFactory.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE)).thenReturn(outputStream1, outputStream2);
    InsightCloseableRegistry closableRegistry = new InsightCloseableRegistry();
    KeyGroupRange keyGroupRange = new KeyGroupRange(0, 2);
    StateSnapshotContextSynchronousImpl context = new StateSnapshotContextSynchronousImpl(checkpointId, checkpointTimestamp, streamFactory, keyGroupRange, closableRegistry);
    // creating the output streams
    context.getRawKeyedOperatorStateOutput();
    context.getRawOperatorStateOutput();
    verify(streamFactory, times(2)).createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);
    assertEquals(2, closableRegistry.size());
    assertTrue(closableRegistry.contains(outputStream1));
    assertTrue(closableRegistry.contains(outputStream2));
    context.closeExceptionally();
    verify(outputStream1).close();
    verify(outputStream2).close();
    assertEquals(0, closableRegistry.size());
}

Example 40 with KeyGroupRange

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