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

Example 46 with KeyGroupRange

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

the class AbstractStreamOperatorTestHarness method repartitionOperatorState.

/**
 * Returns the reshaped the state handles to include only those key-group states in the local
 * key-group range and the operator states that would be assigned to the local subtask.
 */
public static OperatorSubtaskState repartitionOperatorState(final OperatorSubtaskState operatorStateHandles, final int numKeyGroups, final int oldParallelism, final int newParallelism, final int subtaskIndex) {
    Preconditions.checkNotNull(operatorStateHandles, "the previous operatorStateHandles should not be null.");
    // create a new OperatorStateHandles that only contains the state for our key-groups
    List<KeyGroupRange> keyGroupPartitions = StateAssignmentOperation.createKeyGroupPartitions(numKeyGroups, newParallelism);
    KeyGroupRange localKeyGroupRange = keyGroupPartitions.get(subtaskIndex);
    List<KeyedStateHandle> localManagedKeyGroupState = new ArrayList<>();
    StateAssignmentOperation.extractIntersectingState(operatorStateHandles.getManagedKeyedState(), localKeyGroupRange, localManagedKeyGroupState);
    List<KeyedStateHandle> localRawKeyGroupState = new ArrayList<>();
    StateAssignmentOperation.extractIntersectingState(operatorStateHandles.getRawKeyedState(), localKeyGroupRange, localRawKeyGroupState);
    StateObjectCollection<OperatorStateHandle> managedOperatorStates = operatorStateHandles.getManagedOperatorState();
    Collection<OperatorStateHandle> localManagedOperatorState;
    if (!managedOperatorStates.isEmpty()) {
        List<List<OperatorStateHandle>> managedOperatorState = managedOperatorStates.stream().map(Collections::singletonList).collect(Collectors.toList());
        localManagedOperatorState = operatorStateRepartitioner.repartitionState(managedOperatorState, oldParallelism, newParallelism).get(subtaskIndex);
    } else {
        localManagedOperatorState = Collections.emptyList();
    }
    StateObjectCollection<OperatorStateHandle> rawOperatorStates = operatorStateHandles.getRawOperatorState();
    Collection<OperatorStateHandle> localRawOperatorState;
    if (!rawOperatorStates.isEmpty()) {
        List<List<OperatorStateHandle>> rawOperatorState = rawOperatorStates.stream().map(Collections::singletonList).collect(Collectors.toList());
        localRawOperatorState = operatorStateRepartitioner.repartitionState(rawOperatorState, oldParallelism, newParallelism).get(subtaskIndex);
    } else {
        localRawOperatorState = Collections.emptyList();
    }
    return OperatorSubtaskState.builder().setManagedOperatorState(new StateObjectCollection<>(nullToEmptyCollection(localManagedOperatorState))).setRawOperatorState(new StateObjectCollection<>(nullToEmptyCollection(localRawOperatorState))).setManagedKeyedState(new StateObjectCollection<>(nullToEmptyCollection(localManagedKeyGroupState))).setRawKeyedState(new StateObjectCollection<>(nullToEmptyCollection(localRawKeyGroupState))).build();
}

Example 47 with KeyGroupRange

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

the class CheckpointMessagesTest method testConfirmTaskCheckpointed.

@Test
public void testConfirmTaskCheckpointed() {
    final Random rnd = new Random();
    try {
        AcknowledgeCheckpoint noState = new AcknowledgeCheckpoint(new JobID(), new ExecutionAttemptID(), 569345L);
        KeyGroupRange keyGroupRange = KeyGroupRange.of(42, 42);
        TaskStateSnapshot checkpointStateHandles = new TaskStateSnapshot();
        OperatorSubtaskState subtaskState = OperatorSubtaskState.builder().setManagedOperatorState(generatePartitionableStateHandle(new JobVertexID(), 0, 2, 8, false)).setManagedKeyedState(generateKeyGroupState(keyGroupRange, Collections.singletonList(new MyHandle()))).setInputChannelState(singleton(createNewInputChannelStateHandle(10, rnd))).setResultSubpartitionState(singleton(createNewResultSubpartitionStateHandle(10, rnd))).build();
        checkpointStateHandles.putSubtaskStateByOperatorID(new OperatorID(), subtaskState);
        AcknowledgeCheckpoint withState = new AcknowledgeCheckpoint(new JobID(), new ExecutionAttemptID(), 87658976143L, new CheckpointMetrics(), checkpointStateHandles);
        testSerializabilityEqualsHashCode(noState);
        testSerializabilityEqualsHashCode(withState);
    } catch (Exception e) {
        e.printStackTrace();
        fail(e.getMessage());
    }
}

Example 48 with KeyGroupRange

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

the class StateBackendBenchmarkUtils method createRocksDBKeyedStateBackend.

private static RocksDBKeyedStateBackend<Long> createRocksDBKeyedStateBackend(File rootDir) throws IOException {
    File recoveryBaseDir = prepareDirectory(recoveryDirName, rootDir);
    File dbPathFile = prepareDirectory(dbDirName, rootDir);
    ExecutionConfig executionConfig = new ExecutionConfig();
    RocksDBResourceContainer resourceContainer = new RocksDBResourceContainer();
    RocksDBKeyedStateBackendBuilder<Long> builder = new RocksDBKeyedStateBackendBuilder<>("Test", Thread.currentThread().getContextClassLoader(), dbPathFile, resourceContainer, stateName -> resourceContainer.getColumnOptions(), null, LongSerializer.INSTANCE, 2, new KeyGroupRange(0, 1), executionConfig, new LocalRecoveryConfig(null), EmbeddedRocksDBStateBackend.PriorityQueueStateType.ROCKSDB, TtlTimeProvider.DEFAULT, LatencyTrackingStateConfig.disabled(), new UnregisteredMetricsGroup(), Collections.emptyList(), AbstractStateBackend.getCompressionDecorator(executionConfig), new CloseableRegistry());
    try {
        return builder.build();
    } catch (Exception e) {
        IOUtils.closeQuietly(resourceContainer);
        throw e;
    }
}

Example 49 with KeyGroupRange

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

the class StateBackendBenchmarkUtils method createHeapKeyedStateBackend.

private static HeapKeyedStateBackend<Long> createHeapKeyedStateBackend(File rootDir) throws IOException {
    File recoveryBaseDir = prepareDirectory(recoveryDirName, rootDir);
    KeyGroupRange keyGroupRange = new KeyGroupRange(0, 1);
    int numberOfKeyGroups = keyGroupRange.getNumberOfKeyGroups();
    ExecutionConfig executionConfig = new ExecutionConfig();
    HeapPriorityQueueSetFactory priorityQueueSetFactory = new HeapPriorityQueueSetFactory(keyGroupRange, numberOfKeyGroups, 128);
    HeapKeyedStateBackendBuilder<Long> backendBuilder = new HeapKeyedStateBackendBuilder<>(null, new LongSerializer(), Thread.currentThread().getContextClassLoader(), numberOfKeyGroups, keyGroupRange, executionConfig, TtlTimeProvider.DEFAULT, LatencyTrackingStateConfig.disabled(), Collections.emptyList(), AbstractStateBackend.getCompressionDecorator(executionConfig), new LocalRecoveryConfig(null), priorityQueueSetFactory, false, new CloseableRegistry());
    return backendBuilder.build();
}

Example 50 with KeyGroupRange

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

the class RocksIncrementalSnapshotStrategyTest method createSnapshotStrategy.

public RocksIncrementalSnapshotStrategy createSnapshotStrategy(CloseableRegistry closeableRegistry) throws IOException, RocksDBException {
    ColumnFamilyHandle columnFamilyHandle = rocksDBResource.createNewColumnFamily("test");
    RocksDB rocksDB = rocksDBResource.getRocksDB();
    byte[] key = "checkpoint".getBytes();
    byte[] val = "incrementalTest".getBytes();
    rocksDB.put(columnFamilyHandle, key, val);
    // construct RocksIncrementalSnapshotStrategy
    long lastCompletedCheckpointId = -1L;
    ResourceGuard rocksDBResourceGuard = new ResourceGuard();
    SortedMap<Long, Map<StateHandleID, StreamStateHandle>> materializedSstFiles = new TreeMap<>();
    LinkedHashMap<String, RocksDBKeyedStateBackend.RocksDbKvStateInfo> kvStateInformation = new LinkedHashMap<>();
    RocksDBStateUploader rocksDBStateUploader = new RocksDBStateUploader(RocksDBOptions.CHECKPOINT_TRANSFER_THREAD_NUM.defaultValue());
    int keyGroupPrefixBytes = CompositeKeySerializationUtils.computeRequiredBytesInKeyGroupPrefix(2);
    RegisteredKeyValueStateBackendMetaInfo<Integer, ArrayList<Integer>> metaInfo = new RegisteredKeyValueStateBackendMetaInfo<>(StateDescriptor.Type.VALUE, "test", IntSerializer.INSTANCE, new ArrayListSerializer<>(IntSerializer.INSTANCE));
    RocksDBKeyedStateBackend.RocksDbKvStateInfo rocksDbKvStateInfo = new RocksDBKeyedStateBackend.RocksDbKvStateInfo(columnFamilyHandle, metaInfo);
    kvStateInformation.putIfAbsent("test", rocksDbKvStateInfo);
    return new RocksIncrementalSnapshotStrategy<>(rocksDB, rocksDBResourceGuard, IntSerializer.INSTANCE, kvStateInformation, new KeyGroupRange(0, 1), keyGroupPrefixBytes, TestLocalRecoveryConfig.disabled(), closeableRegistry, tmp.newFolder(), UUID.randomUUID(), materializedSstFiles, rocksDBStateUploader, lastCompletedCheckpointId);
}