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

Example 56 with KeyGroupRange

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

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

the class KvStateRequestSerializerTest method testMapSerialization.

/**
	 * Tests map serialization utils.
	 */
@Test
public void testMapSerialization() throws Exception {
    final long key = 0L;
    // objects for heap state list serialisation
    final HeapKeyedStateBackend<Long> longHeapKeyedStateBackend = new HeapKeyedStateBackend<>(mock(TaskKvStateRegistry.class), LongSerializer.INSTANCE, ClassLoader.getSystemClassLoader(), 1, new KeyGroupRange(0, 0), async, new ExecutionConfig());
    longHeapKeyedStateBackend.setCurrentKey(key);
    final InternalMapState<VoidNamespace, Long, String> mapState = (InternalMapState<VoidNamespace, Long, String>) longHeapKeyedStateBackend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, new MapStateDescriptor<>("test", LongSerializer.INSTANCE, StringSerializer.INSTANCE));
    testMapSerialization(key, mapState);
}

Example 58 with KeyGroupRange

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

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

the class AbstractStreamOperator method snapshotState.

@Override
public final OperatorSnapshotResult snapshotState(long checkpointId, long timestamp, CheckpointOptions checkpointOptions) throws Exception {
    KeyGroupRange keyGroupRange = null != keyedStateBackend ? keyedStateBackend.getKeyGroupRange() : KeyGroupRange.EMPTY_KEY_GROUP_RANGE;
    OperatorSnapshotResult snapshotInProgress = new OperatorSnapshotResult();
    CheckpointStreamFactory factory = getCheckpointStreamFactory(checkpointOptions);
    try (StateSnapshotContextSynchronousImpl snapshotContext = new StateSnapshotContextSynchronousImpl(checkpointId, timestamp, factory, keyGroupRange, getContainingTask().getCancelables())) {
        snapshotState(snapshotContext);
        snapshotInProgress.setKeyedStateRawFuture(snapshotContext.getKeyedStateStreamFuture());
        snapshotInProgress.setOperatorStateRawFuture(snapshotContext.getOperatorStateStreamFuture());
        if (null != operatorStateBackend) {
            snapshotInProgress.setOperatorStateManagedFuture(operatorStateBackend.snapshot(checkpointId, timestamp, factory, checkpointOptions));
        }
        if (null != keyedStateBackend) {
            snapshotInProgress.setKeyedStateManagedFuture(keyedStateBackend.snapshot(checkpointId, timestamp, factory, checkpointOptions));
        }
    } catch (Exception snapshotException) {
        try {
            snapshotInProgress.cancel();
        } catch (Exception e) {
            snapshotException.addSuppressed(e);
        }
        throw new Exception("Could not complete snapshot " + checkpointId + " for operator " + getOperatorName() + '.', snapshotException);
    }
    return snapshotInProgress;
}

Example 60 with KeyGroupRange

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

the class HeapInternalTimerServiceTest method testSnapshotAndRebalancingRestore.

/**
	 * This test checks whether timers are assigned to correct key groups
	 * and whether snapshot/restore respects key groups.
	 */
@Test
public void testSnapshotAndRebalancingRestore() throws Exception {
    @SuppressWarnings("unchecked") Triggerable<Integer, String> mockTriggerable = mock(Triggerable.class);
    TestKeyContext keyContext = new TestKeyContext();
    TestProcessingTimeService processingTimeService = new TestProcessingTimeService();
    HeapInternalTimerService<Integer, String> timerService = createTimerService(mockTriggerable, keyContext, processingTimeService, testKeyGroupRange, maxParallelism);
    int midpoint = testKeyGroupRange.getStartKeyGroup() + (testKeyGroupRange.getEndKeyGroup() - testKeyGroupRange.getStartKeyGroup()) / 2;
    // get two sub key-ranges so that we can restore two ranges separately
    KeyGroupRange subKeyGroupRange1 = new KeyGroupRange(testKeyGroupRange.getStartKeyGroup(), midpoint);
    KeyGroupRange subKeyGroupRange2 = new KeyGroupRange(midpoint + 1, testKeyGroupRange.getEndKeyGroup());
    // get two different keys, one per sub range
    int key1 = getKeyInKeyGroupRange(subKeyGroupRange1, maxParallelism);
    int key2 = getKeyInKeyGroupRange(subKeyGroupRange2, maxParallelism);
    keyContext.setCurrentKey(key1);
    timerService.registerProcessingTimeTimer("ciao", 10);
    timerService.registerEventTimeTimer("hello", 10);
    keyContext.setCurrentKey(key2);
    timerService.registerEventTimeTimer("ciao", 10);
    timerService.registerProcessingTimeTimer("hello", 10);
    assertEquals(2, timerService.numProcessingTimeTimers());
    assertEquals(1, timerService.numProcessingTimeTimers("hello"));
    assertEquals(1, timerService.numProcessingTimeTimers("ciao"));
    assertEquals(2, timerService.numEventTimeTimers());
    assertEquals(1, timerService.numEventTimeTimers("hello"));
    assertEquals(1, timerService.numEventTimeTimers("ciao"));
    // one map per sub key-group range
    Map<Integer, byte[]> snapshot1 = new HashMap<>();
    Map<Integer, byte[]> snapshot2 = new HashMap<>();
    for (Integer keyGroupIndex : testKeyGroupRange) {
        ByteArrayOutputStream outStream = new ByteArrayOutputStream();
        timerService.snapshotTimersForKeyGroup(new DataOutputViewStreamWrapper(outStream), keyGroupIndex);
        outStream.close();
        if (subKeyGroupRange1.contains(keyGroupIndex)) {
            snapshot1.put(keyGroupIndex, outStream.toByteArray());
        } else if (subKeyGroupRange2.contains(keyGroupIndex)) {
            snapshot2.put(keyGroupIndex, outStream.toByteArray());
        } else {
            throw new IllegalStateException("Key-Group index doesn't belong to any sub range.");
        }
    }
    // from now on we need everything twice. once per sub key-group range
    @SuppressWarnings("unchecked") Triggerable<Integer, String> mockTriggerable1 = mock(Triggerable.class);
    @SuppressWarnings("unchecked") Triggerable<Integer, String> mockTriggerable2 = mock(Triggerable.class);
    TestKeyContext keyContext1 = new TestKeyContext();
    TestKeyContext keyContext2 = new TestKeyContext();
    TestProcessingTimeService processingTimeService1 = new TestProcessingTimeService();
    TestProcessingTimeService processingTimeService2 = new TestProcessingTimeService();
    HeapInternalTimerService<Integer, String> timerService1 = restoreTimerService(snapshot1, mockTriggerable1, keyContext1, processingTimeService1, subKeyGroupRange1, maxParallelism);
    HeapInternalTimerService<Integer, String> timerService2 = restoreTimerService(snapshot2, mockTriggerable2, keyContext2, processingTimeService2, subKeyGroupRange2, maxParallelism);
    processingTimeService1.setCurrentTime(10);
    timerService1.advanceWatermark(10);
    verify(mockTriggerable1, times(1)).onProcessingTime(anyInternalTimer());
    verify(mockTriggerable1, times(1)).onProcessingTime(eq(new InternalTimer<>(10, key1, "ciao")));
    verify(mockTriggerable1, never()).onProcessingTime(eq(new InternalTimer<>(10, key2, "hello")));
    verify(mockTriggerable1, times(1)).onEventTime(anyInternalTimer());
    verify(mockTriggerable1, times(1)).onEventTime(eq(new InternalTimer<>(10, key1, "hello")));
    verify(mockTriggerable1, never()).onEventTime(eq(new InternalTimer<>(10, key2, "ciao")));
    assertEquals(0, timerService1.numEventTimeTimers());
    processingTimeService2.setCurrentTime(10);
    timerService2.advanceWatermark(10);
    verify(mockTriggerable2, times(1)).onProcessingTime(anyInternalTimer());
    verify(mockTriggerable2, never()).onProcessingTime(eq(new InternalTimer<>(10, key1, "ciao")));
    verify(mockTriggerable2, times(1)).onProcessingTime(eq(new InternalTimer<>(10, key2, "hello")));
    verify(mockTriggerable2, times(1)).onEventTime(anyInternalTimer());
    verify(mockTriggerable2, never()).onEventTime(eq(new InternalTimer<>(10, key1, "hello")));
    verify(mockTriggerable2, times(1)).onEventTime(eq(new InternalTimer<>(10, key2, "ciao")));
    assertEquals(0, timerService2.numEventTimeTimers());
}