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

Example 1 with KeyGroupEntry

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

the class HeapSavepointRestoreOperation method readKeyGroupStateData.

private void readKeyGroupStateData(KeyGroup keyGroup, TypeSerializer<K> keySerializer, Map<Integer, StateMetaInfoSnapshot> kvStatesById) throws Exception {
    try (ThrowingIterator<KeyGroupEntry> entries = keyGroup.getKeyGroupEntries()) {
        while (entries.hasNext()) {
            KeyGroupEntry groupEntry = entries.next();
            StateMetaInfoSnapshot infoSnapshot = kvStatesById.get(groupEntry.getKvStateId());
            switch(infoSnapshot.getBackendStateType()) {
                case KEY_VALUE:
                    readKVStateData(keySerializer, groupEntry, infoSnapshot);
                    break;
                case PRIORITY_QUEUE:
                    readPriorityQueue(groupEntry, infoSnapshot);
                    break;
                case OPERATOR:
                case BROADCAST:
                    throw new IllegalStateException("Expected only keyed state. Received: " + infoSnapshot.getBackendStateType());
            }
        }
    }
}

Example 2 with KeyGroupEntry

use of org.apache.flink.runtime.state.restore.KeyGroupEntry 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 3 with KeyGroupEntry

use of org.apache.flink.runtime.state.restore.KeyGroupEntry 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());
                }
            }
        }
    }
}