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

Example 86 with KeyGroupRange

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

the class AbstractStreamOperatorTest method testStateAndTimerStateShufflingScalingUp.

/**
 * Verify that state and timers are checkpointed per key group and that they are correctly
 * assigned to operator subtasks when restoring.
 */
@Test
public void testStateAndTimerStateShufflingScalingUp() throws Exception {
    final int maxParallelism = 10;
    // first get two keys that will fall into different key-group ranges that go
    // to different operator subtasks when we restore
    // get two sub key-ranges so that we can restore two ranges separately
    KeyGroupRange subKeyGroupRange1 = new KeyGroupRange(0, (maxParallelism / 2) - 1);
    KeyGroupRange subKeyGroupRange2 = new KeyGroupRange(subKeyGroupRange1.getEndKeyGroup() + 1, maxParallelism - 1);
    // get two different keys, one per sub range
    int key1 = getKeyInKeyGroupRange(subKeyGroupRange1, maxParallelism);
    int key2 = getKeyInKeyGroupRange(subKeyGroupRange2, maxParallelism);
    OperatorSubtaskState snapshot;
    try (KeyedOneInputStreamOperatorTestHarness<Integer, Tuple2<Integer, String>, String> testHarness = createTestHarness(maxParallelism, 1, 0)) {
        testHarness.open();
        testHarness.processWatermark(0L);
        testHarness.setProcessingTime(0L);
        testHarness.processElement(new Tuple2<>(key1, "SET_EVENT_TIME_TIMER:10"), 0);
        testHarness.processElement(new Tuple2<>(key2, "SET_EVENT_TIME_TIMER:20"), 0);
        testHarness.processElement(new Tuple2<>(key1, "SET_PROC_TIME_TIMER:10"), 0);
        testHarness.processElement(new Tuple2<>(key2, "SET_PROC_TIME_TIMER:20"), 0);
        testHarness.processElement(new Tuple2<>(key1, "SET_STATE:HELLO"), 0);
        testHarness.processElement(new Tuple2<>(key2, "SET_STATE:CIAO"), 0);
        assertTrue(extractResult(testHarness).isEmpty());
        snapshot = testHarness.snapshot(0, 0);
    }
    // now, restore in two operators, first operator 1
    OperatorSubtaskState initState1 = AbstractStreamOperatorTestHarness.repartitionOperatorState(snapshot, maxParallelism, 1, 2, 0);
    try (KeyedOneInputStreamOperatorTestHarness<Integer, Tuple2<Integer, String>, String> testHarness1 = createTestHarness(maxParallelism, 2, 0)) {
        testHarness1.setup();
        testHarness1.initializeState(initState1);
        testHarness1.open();
        testHarness1.processWatermark(10L);
        assertThat(extractResult(testHarness1), contains("ON_EVENT_TIME:HELLO"));
        assertTrue(extractResult(testHarness1).isEmpty());
        // this should not trigger anything, the trigger for WM=20 should sit in the
        // other operator subtask
        testHarness1.processWatermark(20L);
        assertTrue(extractResult(testHarness1).isEmpty());
        testHarness1.setProcessingTime(10L);
        assertThat(extractResult(testHarness1), contains("ON_PROC_TIME:HELLO"));
        assertTrue(extractResult(testHarness1).isEmpty());
        // this should not trigger anything, the trigger for TIME=20 should sit in the
        // other operator subtask
        testHarness1.setProcessingTime(20L);
        assertTrue(extractResult(testHarness1).isEmpty());
    }
    // now, for the second operator
    OperatorSubtaskState initState2 = AbstractStreamOperatorTestHarness.repartitionOperatorState(snapshot, maxParallelism, 1, 2, 1);
    try (KeyedOneInputStreamOperatorTestHarness<Integer, Tuple2<Integer, String>, String> testHarness2 = createTestHarness(maxParallelism, 2, 1)) {
        testHarness2.setup();
        testHarness2.initializeState(initState2);
        testHarness2.open();
        testHarness2.processWatermark(10L);
        // nothing should happen because this timer is in the other subtask
        assertTrue(extractResult(testHarness2).isEmpty());
        testHarness2.processWatermark(20L);
        assertThat(extractResult(testHarness2), contains("ON_EVENT_TIME:CIAO"));
        testHarness2.setProcessingTime(10L);
        // nothing should happen because this timer is in the other subtask
        assertTrue(extractResult(testHarness2).isEmpty());
        testHarness2.setProcessingTime(20L);
        assertThat(extractResult(testHarness2), contains("ON_PROC_TIME:CIAO"));
        assertTrue(extractResult(testHarness2).isEmpty());
    }
}

Example 87 with KeyGroupRange

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

the class AbstractStreamOperatorTest method testStateAndTimerStateShufflingScalingDown.

@Test
public void testStateAndTimerStateShufflingScalingDown() throws Exception {
    final int maxParallelism = 10;
    // first get two keys that will fall into different key-group ranges that go
    // to different operator subtasks when we restore
    // get two sub key-ranges so that we can restore two ranges separately
    KeyGroupRange subKeyGroupRange1 = new KeyGroupRange(0, (maxParallelism / 2) - 1);
    KeyGroupRange subKeyGroupRange2 = new KeyGroupRange(subKeyGroupRange1.getEndKeyGroup() + 1, maxParallelism - 1);
    // get two different keys, one per sub range
    int key1 = getKeyInKeyGroupRange(subKeyGroupRange1, maxParallelism);
    int key2 = getKeyInKeyGroupRange(subKeyGroupRange2, maxParallelism);
    OperatorSubtaskState snapshot1, snapshot2;
    // register some state with both instances and scale down to parallelism 1
    try (KeyedOneInputStreamOperatorTestHarness<Integer, Tuple2<Integer, String>, String> testHarness1 = createTestHarness(maxParallelism, 2, 0)) {
        testHarness1.setup();
        testHarness1.open();
        testHarness1.processWatermark(0L);
        testHarness1.setProcessingTime(0L);
        testHarness1.processElement(new Tuple2<>(key1, "SET_EVENT_TIME_TIMER:30"), 0);
        testHarness1.processElement(new Tuple2<>(key1, "SET_PROC_TIME_TIMER:30"), 0);
        testHarness1.processElement(new Tuple2<>(key1, "SET_STATE:HELLO"), 0);
        snapshot1 = testHarness1.snapshot(0, 0);
    }
    try (KeyedOneInputStreamOperatorTestHarness<Integer, Tuple2<Integer, String>, String> testHarness2 = createTestHarness(maxParallelism, 2, 1)) {
        testHarness2.setup();
        testHarness2.open();
        testHarness2.processWatermark(0L);
        testHarness2.setProcessingTime(0L);
        testHarness2.processElement(new Tuple2<>(key2, "SET_EVENT_TIME_TIMER:40"), 0);
        testHarness2.processElement(new Tuple2<>(key2, "SET_PROC_TIME_TIMER:40"), 0);
        testHarness2.processElement(new Tuple2<>(key2, "SET_STATE:CIAO"), 0);
        snapshot2 = testHarness2.snapshot(0, 0);
    }
    // take a snapshot from each one of the "parallel" instances of the operator
    // and combine them into one so that we can scale down
    OperatorSubtaskState repackagedState = AbstractStreamOperatorTestHarness.repackageState(snapshot1, snapshot2);
    OperatorSubtaskState initSubTaskState = AbstractStreamOperatorTestHarness.repartitionOperatorState(repackagedState, maxParallelism, 2, 1, 0);
    try (KeyedOneInputStreamOperatorTestHarness<Integer, Tuple2<Integer, String>, String> testHarness3 = createTestHarness(maxParallelism, 1, 0)) {
        testHarness3.setup();
        testHarness3.initializeState(initSubTaskState);
        testHarness3.open();
        testHarness3.processWatermark(30L);
        assertThat(extractResult(testHarness3), contains("ON_EVENT_TIME:HELLO"));
        assertTrue(extractResult(testHarness3).isEmpty());
        testHarness3.processWatermark(40L);
        assertThat(extractResult(testHarness3), contains("ON_EVENT_TIME:CIAO"));
        assertTrue(extractResult(testHarness3).isEmpty());
        testHarness3.setProcessingTime(30L);
        assertThat(extractResult(testHarness3), contains("ON_PROC_TIME:HELLO"));
        assertTrue(extractResult(testHarness3).isEmpty());
        testHarness3.setProcessingTime(40L);
        assertThat(extractResult(testHarness3), contains("ON_PROC_TIME:CIAO"));
        assertTrue(extractResult(testHarness3).isEmpty());
    }
}

Example 88 with KeyGroupRange

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

the class OperatorSubtaskStateTest method testToBuilderCorrectness.

@Test
public void testToBuilderCorrectness() throws IOException {
    // given: Initialized operator subtask state.
    JobVertexID jobVertexID = new JobVertexID();
    int index = 0;
    Random random = new Random();
    OperatorSubtaskState operatorSubtaskState = OperatorSubtaskState.builder().setManagedOperatorState(generatePartitionableStateHandle(jobVertexID, index, 2, 8, false)).setRawOperatorState(generatePartitionableStateHandle(jobVertexID, index, 2, 8, true)).setManagedKeyedState(generateKeyGroupState(jobVertexID, new KeyGroupRange(0, 11), false)).setRawKeyedState(generateKeyGroupState(jobVertexID, new KeyGroupRange(0, 9), true)).setInputChannelState(StateObjectCollection.singleton(createNewInputChannelStateHandle(3, random))).setResultSubpartitionState(StateObjectCollection.singleton(createNewResultSubpartitionStateHandle(3, random))).setInputRescalingDescriptor(InflightDataRescalingDescriptorUtil.rescalingDescriptor(new int[1], new RescaleMappings[0], Collections.singleton(1))).setOutputRescalingDescriptor(InflightDataRescalingDescriptorUtil.rescalingDescriptor(new int[1], new RescaleMappings[0], Collections.singleton(2))).build();
    // when: Copy the operator subtask state.
    OperatorSubtaskState operatorSubtaskStateCopy = operatorSubtaskState.toBuilder().build();
    // then: It should be equal to original one.
    assertTrue(reflectionEquals(operatorSubtaskState, operatorSubtaskStateCopy));
}

Example 89 with KeyGroupRange

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

the class CheckpointStateRestoreTest method testSetState.

/**
 * Tests that on restore the task state is reset for each stateful task.
 */
@Test
public void testSetState() {
    try {
        KeyGroupRange keyGroupRange = KeyGroupRange.of(0, 0);
        List<SerializableObject> testStates = Collections.singletonList(new SerializableObject());
        final KeyedStateHandle serializedKeyGroupStates = CheckpointCoordinatorTestingUtils.generateKeyGroupState(keyGroupRange, testStates);
        final JobVertexID statefulId = new JobVertexID();
        final JobVertexID statelessId = new JobVertexID();
        ExecutionGraph graph = new CheckpointCoordinatorTestingUtils.CheckpointExecutionGraphBuilder().addJobVertex(statefulId, 3, 256).addJobVertex(statelessId, 2, 256).build();
        ExecutionJobVertex stateful = graph.getJobVertex(statefulId);
        ExecutionJobVertex stateless = graph.getJobVertex(statelessId);
        ExecutionVertex stateful1 = stateful.getTaskVertices()[0];
        ExecutionVertex stateful2 = stateful.getTaskVertices()[1];
        ExecutionVertex stateful3 = stateful.getTaskVertices()[2];
        ExecutionVertex stateless1 = stateless.getTaskVertices()[0];
        ExecutionVertex stateless2 = stateless.getTaskVertices()[1];
        Execution statefulExec1 = stateful1.getCurrentExecutionAttempt();
        Execution statefulExec2 = stateful2.getCurrentExecutionAttempt();
        Execution statefulExec3 = stateful3.getCurrentExecutionAttempt();
        Execution statelessExec1 = stateless1.getCurrentExecutionAttempt();
        Execution statelessExec2 = stateless2.getCurrentExecutionAttempt();
        ManuallyTriggeredScheduledExecutor manuallyTriggeredScheduledExecutor = new ManuallyTriggeredScheduledExecutor();
        CheckpointCoordinator coord = new CheckpointCoordinatorBuilder().setExecutionGraph(graph).setTimer(manuallyTriggeredScheduledExecutor).build();
        // create ourselves a checkpoint with state
        coord.triggerCheckpoint(false);
        manuallyTriggeredScheduledExecutor.triggerAll();
        PendingCheckpoint pending = coord.getPendingCheckpoints().values().iterator().next();
        final long checkpointId = pending.getCheckpointId();
        final TaskStateSnapshot subtaskStates = new TaskStateSnapshot();
        subtaskStates.putSubtaskStateByOperatorID(OperatorID.fromJobVertexID(statefulId), OperatorSubtaskState.builder().setManagedKeyedState(serializedKeyGroupStates).build());
        coord.receiveAcknowledgeMessage(new AcknowledgeCheckpoint(graph.getJobID(), statefulExec1.getAttemptId(), checkpointId, new CheckpointMetrics(), subtaskStates), TASK_MANAGER_LOCATION_INFO);
        coord.receiveAcknowledgeMessage(new AcknowledgeCheckpoint(graph.getJobID(), statefulExec2.getAttemptId(), checkpointId, new CheckpointMetrics(), subtaskStates), TASK_MANAGER_LOCATION_INFO);
        coord.receiveAcknowledgeMessage(new AcknowledgeCheckpoint(graph.getJobID(), statefulExec3.getAttemptId(), checkpointId, new CheckpointMetrics(), subtaskStates), TASK_MANAGER_LOCATION_INFO);
        coord.receiveAcknowledgeMessage(new AcknowledgeCheckpoint(graph.getJobID(), statelessExec1.getAttemptId(), checkpointId), TASK_MANAGER_LOCATION_INFO);
        coord.receiveAcknowledgeMessage(new AcknowledgeCheckpoint(graph.getJobID(), statelessExec2.getAttemptId(), checkpointId), TASK_MANAGER_LOCATION_INFO);
        assertEquals(1, coord.getNumberOfRetainedSuccessfulCheckpoints());
        assertEquals(0, coord.getNumberOfPendingCheckpoints());
        // let the coordinator inject the state
        assertTrue(coord.restoreLatestCheckpointedStateToAll(new HashSet<>(Arrays.asList(stateful, stateless)), false));
        // verify that each stateful vertex got the state
        assertEquals(subtaskStates, statefulExec1.getTaskRestore().getTaskStateSnapshot());
        assertEquals(subtaskStates, statefulExec2.getTaskRestore().getTaskStateSnapshot());
        assertEquals(subtaskStates, statefulExec3.getTaskRestore().getTaskStateSnapshot());
        assertNull(statelessExec1.getTaskRestore());
        assertNull(statelessExec2.getTaskRestore());
    } catch (Exception e) {
        e.printStackTrace();
        fail(e.getMessage());
    }
}

Example 90 with KeyGroupRange

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

the class PrioritizedOperatorSubtaskStateTest method generateForConfiguration.

/**
 * Generator for all 3^4 = 81 possible configurations of a OperatorSubtaskState: - 4 different
 * sub-states: managed/raw + operator/keyed. - 3 different options per sub-state: empty
 * (simulate no state), single handle (simulate recovery), 2 handles (simulate e.g. rescaling)
 */
private OperatorSubtaskState generateForConfiguration(int conf) {
    // 3^4
    Preconditions.checkState(conf >= 0 && conf <= 80);
    final int numModes = 3;
    KeyGroupRange keyGroupRange = new KeyGroupRange(0, 4);
    KeyGroupRange keyGroupRange1 = new KeyGroupRange(0, 2);
    KeyGroupRange keyGroupRange2 = new KeyGroupRange(3, 4);
    int div = 1;
    int mode = (conf / div) % numModes;
    StateObjectCollection<OperatorStateHandle> s1 = mode == 0 ? StateObjectCollection.empty() : mode == 1 ? new StateObjectCollection<>(Collections.singletonList(createNewOperatorStateHandle(2, RANDOM))) : new StateObjectCollection<>(Arrays.asList(createNewOperatorStateHandle(2, RANDOM), createNewOperatorStateHandle(2, RANDOM)));
    div *= numModes;
    mode = (conf / div) % numModes;
    StateObjectCollection<OperatorStateHandle> s2 = mode == 0 ? StateObjectCollection.empty() : mode == 1 ? new StateObjectCollection<>(Collections.singletonList(createNewOperatorStateHandle(2, RANDOM))) : new StateObjectCollection<>(Arrays.asList(createNewOperatorStateHandle(2, RANDOM), createNewOperatorStateHandle(2, RANDOM)));
    div *= numModes;
    mode = (conf / div) % numModes;
    StateObjectCollection<KeyedStateHandle> s3 = mode == 0 ? StateObjectCollection.empty() : mode == 1 ? new StateObjectCollection<>(Collections.singletonList(createNewKeyedStateHandle(keyGroupRange))) : new StateObjectCollection<>(Arrays.asList(createNewKeyedStateHandle(keyGroupRange1), createNewKeyedStateHandle(keyGroupRange2)));
    div *= numModes;
    mode = (conf / div) % numModes;
    StateObjectCollection<KeyedStateHandle> s4 = mode == 0 ? StateObjectCollection.empty() : mode == 1 ? new StateObjectCollection<>(Collections.singletonList(createNewKeyedStateHandle(keyGroupRange))) : new StateObjectCollection<>(Arrays.asList(createNewKeyedStateHandle(keyGroupRange1), createNewKeyedStateHandle(keyGroupRange2)));
    return OperatorSubtaskState.builder().setManagedOperatorState(s1).setRawOperatorState(s2).setManagedKeyedState(s3).setRawKeyedState(s4).build();
}