Github user sihuazhou commented on a diff in the pull request: https://github.com/apache/flink/pull/5582#discussion_r191447222 --- Diff: flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java --- @@ -718,70 +718,63 @@ void restore(Collection<KeyedStateHandle> restoreStateHandles) throws Exception boolean hasExtraKeys = (restoreStateHandles.size() > 1 || !Objects.equals(restoreStateHandles.iterator().next().getKeyGroupRange(), stateBackend.keyGroupRange)); - if (hasExtraKeys) { - stateBackend.createDB(); - } - - for (KeyedStateHandle rawStateHandle : restoreStateHandles) { - - if (rawStateHandle instanceof IncrementalKeyedStateHandle) { - restoreInstance((IncrementalKeyedStateHandle) rawStateHandle, hasExtraKeys); - } else if (rawStateHandle instanceof IncrementalLocalKeyedStateHandle) { - Preconditions.checkState(!hasExtraKeys, "Cannot recover from local state after rescaling."); - restoreInstance((IncrementalLocalKeyedStateHandle) rawStateHandle); - } else { - throw new IllegalStateException("Unexpected state handle type, " + - "expected " + IncrementalKeyedStateHandle.class + - ", but found " + rawStateHandle.getClass()); - } + if (!hasExtraKeys) { + restoreFromSingleHandle(restoreStateHandles.iterator().next()); --- End diff -- You are right, I think this should be improved.
---