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

Example 1 with KeyGroup

use of org.apache.flink.runtime.state.restore.KeyGroup in project flink by apache.

the class RocksDBHeapTimersFullRestoreOperation method applyRestoreResult.

private void applyRestoreResult(SavepointRestoreResult savepointRestoreResult) throws IOException, RocksDBException, StateMigrationException {
    List<StateMetaInfoSnapshot> restoredMetaInfos = savepointRestoreResult.getStateMetaInfoSnapshots();
    Map<Integer, ColumnFamilyHandle> columnFamilyHandles = new HashMap<>();
    Map<Integer, HeapPriorityQueueSnapshotRestoreWrapper<?>> restoredPQStates = new HashMap<>();
    for (int i = 0; i < restoredMetaInfos.size(); i++) {
        StateMetaInfoSnapshot restoredMetaInfo = restoredMetaInfos.get(i);
        if (restoredMetaInfo.getBackendStateType() == BackendStateType.PRIORITY_QUEUE) {
            String stateName = restoredMetaInfo.getName();
            HeapPriorityQueueSnapshotRestoreWrapper<?> queueWrapper = registeredPQStates.computeIfAbsent(stateName, key -> createInternal(new RegisteredPriorityQueueStateBackendMetaInfo<>(restoredMetaInfo)));
            restoredPQStates.put(i, queueWrapper);
        } else {
            RocksDbKvStateInfo registeredStateCFHandle = this.rocksHandle.getOrRegisterStateColumnFamilyHandle(null, restoredMetaInfo);
            columnFamilyHandles.put(i, registeredStateCFHandle.columnFamilyHandle);
        }
    }
    try (ThrowingIterator<KeyGroup> keyGroups = savepointRestoreResult.getRestoredKeyGroups()) {
        restoreKVStateData(keyGroups, columnFamilyHandles, restoredPQStates);
    }
}

Example 2 with KeyGroup

use of org.apache.flink.runtime.state.restore.KeyGroup in project flink by apache.

the class RocksDBFullRestoreOperation method applyRestoreResult.

private void applyRestoreResult(SavepointRestoreResult savepointRestoreResult) throws IOException, RocksDBException, StateMigrationException {
    List<StateMetaInfoSnapshot> restoredMetaInfos = savepointRestoreResult.getStateMetaInfoSnapshots();
    Map<Integer, ColumnFamilyHandle> columnFamilyHandles = new HashMap<>();
    for (int i = 0; i < restoredMetaInfos.size(); i++) {
        StateMetaInfoSnapshot restoredMetaInfo = restoredMetaInfos.get(i);
        RocksDbKvStateInfo registeredStateCFHandle = this.rocksHandle.getOrRegisterStateColumnFamilyHandle(null, restoredMetaInfo);
        columnFamilyHandles.put(i, registeredStateCFHandle.columnFamilyHandle);
    }
    try (ThrowingIterator<KeyGroup> keyGroups = savepointRestoreResult.getRestoredKeyGroups()) {
        restoreKVStateData(keyGroups, columnFamilyHandles);
    }
}

Example 3 with KeyGroup

use of org.apache.flink.runtime.state.restore.KeyGroup in project flink by apache.

the class HeapSavepointRestoreOperation method restore.

@Override
public Void restore() throws Exception {
    registeredKVStates.clear();
    registeredPQStates.clear();
    try (ThrowingIterator<SavepointRestoreResult> restore = this.savepointRestoreOperation.restore()) {
        while (restore.hasNext()) {
            SavepointRestoreResult restoreResult = restore.next();
            List<StateMetaInfoSnapshot> restoredMetaInfos = restoreResult.getStateMetaInfoSnapshots();
            final Map<Integer, StateMetaInfoSnapshot> kvStatesById = this.heapMetaInfoRestoreOperation.createOrCheckStateForMetaInfo(restoredMetaInfos, registeredKVStates, registeredPQStates);
            try (ThrowingIterator<KeyGroup> keyGroups = restoreResult.getRestoredKeyGroups()) {
                while (keyGroups.hasNext()) {
                    readKeyGroupStateData(keyGroups.next(), keySerializerProvider.previousSchemaSerializer(), kvStatesById);
                }
            }
        }
    }
    return null;
}

Example 4 with KeyGroup

use of org.apache.flink.runtime.state.restore.KeyGroup in project flink by apache.

the class RocksDBHeapTimersFullRestoreOperation method restoreKVStateData.

/**
 * Restore the KV-state / ColumnFamily data for all key-groups referenced by the current state
 * handle.
 */
private void restoreKVStateData(ThrowingIterator<KeyGroup> keyGroups, Map<Integer, ColumnFamilyHandle> columnFamilies, Map<Integer, HeapPriorityQueueSnapshotRestoreWrapper<?>> restoredPQStates) throws IOException, RocksDBException, StateMigrationException {
    // for all key-groups in the current state handle...
    try (RocksDBWriteBatchWrapper writeBatchWrapper = new RocksDBWriteBatchWrapper(this.rocksHandle.getDb(), writeBatchSize)) {
        HeapPriorityQueueSnapshotRestoreWrapper<HeapPriorityQueueElement> restoredPQ = null;
        ColumnFamilyHandle handle = null;
        while (keyGroups.hasNext()) {
            KeyGroup keyGroup = keyGroups.next();
            try (ThrowingIterator<KeyGroupEntry> groupEntries = keyGroup.getKeyGroupEntries()) {
                int oldKvStateId = -1;
                while (groupEntries.hasNext()) {
                    KeyGroupEntry groupEntry = groupEntries.next();
                    int kvStateId = groupEntry.getKvStateId();
                    if (kvStateId != oldKvStateId) {
                        oldKvStateId = kvStateId;
                        handle = columnFamilies.get(kvStateId);
                        restoredPQ = getRestoredPQ(restoredPQStates, kvStateId);
                    }
                    if (restoredPQ != null) {
                        restoreQueueElement(restoredPQ, groupEntry);
                    } else if (handle != null) {
                        writeBatchWrapper.put(handle, groupEntry.getKey(), groupEntry.getValue());
                    } else {
                        throw new IllegalStateException("Unknown state id: " + kvStateId);
                    }
                }
            }
        }
    }
}

Example 5 with KeyGroup

use of org.apache.flink.runtime.state.restore.KeyGroup in project flink by apache.

the class RocksDBFullRestoreOperation method restoreKVStateData.

/**
 * Restore the KV-state / ColumnFamily data for all key-groups referenced by the current state
 * handle.
 */
private void restoreKVStateData(ThrowingIterator<KeyGroup> keyGroups, Map<Integer, ColumnFamilyHandle> columnFamilies) throws IOException, RocksDBException, StateMigrationException {
    // for all key-groups in the current state handle...
    try (RocksDBWriteBatchWrapper writeBatchWrapper = new RocksDBWriteBatchWrapper(this.rocksHandle.getDb(), writeBatchSize)) {
        ColumnFamilyHandle handle = null;
        while (keyGroups.hasNext()) {
            KeyGroup keyGroup = keyGroups.next();
            try (ThrowingIterator<KeyGroupEntry> groupEntries = keyGroup.getKeyGroupEntries()) {
                int oldKvStateId = -1;
                while (groupEntries.hasNext()) {
                    KeyGroupEntry groupEntry = groupEntries.next();
                    int kvStateId = groupEntry.getKvStateId();
                    if (kvStateId != oldKvStateId) {
                        oldKvStateId = kvStateId;
                        handle = columnFamilies.get(kvStateId);
                    }
                    writeBatchWrapper.put(handle, groupEntry.getKey(), groupEntry.getValue());
                }
            }
        }
    }
}