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

Example 1 with StateSnapshotRestore

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

the class HeapKeyedStateBackend method getKeys.

@SuppressWarnings("unchecked")
@Override
public <N> Stream<K> getKeys(String state, N namespace) {
    if (!registeredKVStates.containsKey(state)) {
        return Stream.empty();
    }
    final StateSnapshotRestore stateSnapshotRestore = registeredKVStates.get(state);
    StateTable<K, N, ?> table = (StateTable<K, N, ?>) stateSnapshotRestore;
    return table.getKeys(namespace);
}

Example 2 with StateSnapshotRestore

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

the class HeapMetaInfoRestoreOperation method createOrCheckStateForMetaInfo.

Map<Integer, StateMetaInfoSnapshot> createOrCheckStateForMetaInfo(List<StateMetaInfoSnapshot> restoredMetaInfo, Map<String, StateTable<K, ?, ?>> registeredKVStates, Map<String, HeapPriorityQueueSnapshotRestoreWrapper<?>> registeredPQStates) {
    final Map<Integer, StateMetaInfoSnapshot> kvStatesById = new HashMap<>();
    for (StateMetaInfoSnapshot metaInfoSnapshot : restoredMetaInfo) {
        final StateSnapshotRestore registeredState;
        switch(metaInfoSnapshot.getBackendStateType()) {
            case KEY_VALUE:
                registeredState = registeredKVStates.get(metaInfoSnapshot.getName());
                if (registeredState == null) {
                    RegisteredKeyValueStateBackendMetaInfo<?, ?> registeredKeyedBackendStateMetaInfo = new RegisteredKeyValueStateBackendMetaInfo<>(metaInfoSnapshot);
                    registeredKVStates.put(metaInfoSnapshot.getName(), stateTableFactory.newStateTable(keyContext, registeredKeyedBackendStateMetaInfo, keySerializerProvider.currentSchemaSerializer()));
                }
                break;
            case PRIORITY_QUEUE:
                registeredState = registeredPQStates.get(metaInfoSnapshot.getName());
                if (registeredState == null) {
                    registeredPQStates.put(metaInfoSnapshot.getName(), createInternal(new RegisteredPriorityQueueStateBackendMetaInfo<>(metaInfoSnapshot)));
                }
                break;
            default:
                throw new IllegalStateException("Unexpected state type: " + metaInfoSnapshot.getBackendStateType() + ".");
        }
        // always put metaInfo into kvStatesById, because kvStatesById is KeyGroupsStateHandle
        // related
        kvStatesById.put(kvStatesById.size(), metaInfoSnapshot);
    }
    return kvStatesById;
}

Example 3 with StateSnapshotRestore

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

the class HeapRestoreOperation method readKeyGroupStateData.

private void readKeyGroupStateData(InputStream inputStream, Map<Integer, StateMetaInfoSnapshot> kvStatesById, int keyGroupIndex, int numStates, int readVersion) throws IOException {
    DataInputViewStreamWrapper inView = new DataInputViewStreamWrapper(inputStream);
    for (int i = 0; i < numStates; i++) {
        final int kvStateId = inView.readShort();
        final StateMetaInfoSnapshot stateMetaInfoSnapshot = kvStatesById.get(kvStateId);
        final StateSnapshotRestore registeredState;
        switch(stateMetaInfoSnapshot.getBackendStateType()) {
            case KEY_VALUE:
                registeredState = registeredKVStates.get(stateMetaInfoSnapshot.getName());
                break;
            case PRIORITY_QUEUE:
                registeredState = registeredPQStates.get(stateMetaInfoSnapshot.getName());
                break;
            default:
                throw new IllegalStateException("Unexpected state type: " + stateMetaInfoSnapshot.getBackendStateType() + ".");
        }
        StateSnapshotKeyGroupReader keyGroupReader = registeredState.keyGroupReader(readVersion);
        keyGroupReader.readMappingsInKeyGroup(inView, keyGroupIndex);
    }
}

Example 4 with StateSnapshotRestore

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

the class HeapSnapshotResources method processSnapshotMetaInfoForAllStates.

private static void processSnapshotMetaInfoForAllStates(List<StateMetaInfoSnapshot> metaInfoSnapshots, Map<StateUID, StateSnapshot> cowStateStableSnapshots, Map<StateUID, Integer> stateNamesToId, Map<String, ? extends StateSnapshotRestore> registeredStates, StateMetaInfoSnapshot.BackendStateType stateType) {
    for (Map.Entry<String, ? extends StateSnapshotRestore> kvState : registeredStates.entrySet()) {
        final StateUID stateUid = StateUID.of(kvState.getKey(), stateType);
        stateNamesToId.put(stateUid, stateNamesToId.size());
        StateSnapshotRestore state = kvState.getValue();
        if (null != state) {
            final StateSnapshot stateSnapshot = state.stateSnapshot();
            metaInfoSnapshots.add(stateSnapshot.getMetaInfoSnapshot());
            cowStateStableSnapshots.put(stateUid, stateSnapshot);
        }
    }
}

Example 5 with StateSnapshotRestore

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

the class HeapKeyedStateBackend method getKeysAndNamespaces.

@SuppressWarnings("unchecked")
@Override
public <N> Stream<Tuple2<K, N>> getKeysAndNamespaces(String state) {
    if (!registeredKVStates.containsKey(state)) {
        return Stream.empty();
    }
    final StateSnapshotRestore stateSnapshotRestore = registeredKVStates.get(state);
    StateTable<K, N, ?> table = (StateTable<K, N, ?>) stateSnapshotRestore;
    return table.getKeysAndNamespaces();
}