Search in sources :

Example 1 with HeapPriorityQueueElement

use of org.apache.flink.runtime.state.heap.HeapPriorityQueueElement in project flink by apache.

the class RocksDBHeapTimersFullRestoreOperation method restoreQueueElement.

private void restoreQueueElement(HeapPriorityQueueSnapshotRestoreWrapper<HeapPriorityQueueElement> restoredPQ, KeyGroupEntry groupEntry) throws IOException {
    deserializer.setBuffer(groupEntry.getKey());
    deserializer.skipBytesToRead(keyGroupPrefixBytes);
    HeapPriorityQueueElement queueElement = restoredPQ.getMetaInfo().getElementSerializer().deserialize(deserializer);
    restoredPQ.getPriorityQueue().add(queueElement);
}
Also used : HeapPriorityQueueElement(org.apache.flink.runtime.state.heap.HeapPriorityQueueElement)

Example 2 with HeapPriorityQueueElement

use of org.apache.flink.runtime.state.heap.HeapPriorityQueueElement 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);
                    }
                }
            }
        }
    }
}
Also used : KeyGroupEntry(org.apache.flink.runtime.state.restore.KeyGroupEntry) KeyGroup(org.apache.flink.runtime.state.restore.KeyGroup) RocksDBWriteBatchWrapper(org.apache.flink.contrib.streaming.state.RocksDBWriteBatchWrapper) HeapPriorityQueueElement(org.apache.flink.runtime.state.heap.HeapPriorityQueueElement) ColumnFamilyHandle(org.rocksdb.ColumnFamilyHandle)

Aggregations

HeapPriorityQueueElement (org.apache.flink.runtime.state.heap.HeapPriorityQueueElement)2 RocksDBWriteBatchWrapper (org.apache.flink.contrib.streaming.state.RocksDBWriteBatchWrapper)1 KeyGroup (org.apache.flink.runtime.state.restore.KeyGroup)1 KeyGroupEntry (org.apache.flink.runtime.state.restore.KeyGroupEntry)1 ColumnFamilyHandle (org.rocksdb.ColumnFamilyHandle)1