Repository: apache/flink
File: flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapKeyedStateBackend.java (Java)

Code before:
                        + " config has been changed from {} to {}.",
                stateDesc,
                TtlAwareSerializer.isSerializerTtlEnabled(previousSerializer),
                TtlAwareSerializer.isSerializerTtlEnabled(currentSerializer));

        // we need to get an actual state instance because migration is different
        // for different state types. For example, ListState needs to deal with
        // individual elements
        StateCreateFactory stateCreateFactory = STATE_CREATE_FACTORIES.get(stateDesc.getType());
        if (stateCreateFactory == null) {
            throw new FlinkRuntimeException(stateNotSupportedMessage(stateDesc));
        }
        State state =
                stateCreateFactory.createState(stateDesc, stateTable, stateTable.keySerializer);
        if (!(state instanceof AbstractHeapState)) {
            throw new FlinkRuntimeException(
                    "State should be an AbstractRocksDBState but is " + state);
        }
        AbstractHeapState<K, N, V> heapState = (AbstractHeapState<K, N, V>) state;
        TtlAwareSerializer<V, ?> currentTtlAwareSerializer =
                (TtlAwareSerializer<V, ?>)
                        TtlAwareSerializer.wrapTtlAwareSerializer(currentSerializer);

        while (iterator.hasNext()) {
            final StateEntry<K, N, V> entry = iterator.next();
            stateTable.put(
                    entry.getKey(),
                    KeyGroupRangeAssignment.assignToKeyGroup(entry.getKey(), numberOfKeyGroups),
                    entry.getNamespace(),
                    heapState.migrateTtlValue(
                            entry.getState(), currentTtlAwareSerializer, ttlTimeProvider));
        }
    }

    @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);
    }

    @SuppressWarnings("unchecked")
    @Override
    public <N> Stream<K> getKeys(List<String> states, N namespace) {
        final List<StateTable<K, N, ?>> tables =
                states.stream()

Review comment (suggestion):
It there any concurrency issue when performing `stateTable.put()` over the iteration?
If so, I'd suggest a new inner transform method to in-place update all the value of inner entries

Code after:
                        + " config has been changed from {} to {}.",
                stateDesc,
                TtlAwareSerializer.isSerializerTtlEnabled(previousSerializer),
                TtlAwareSerializer.isSerializerTtlEnabled(currentSerializer));

        // we need to get an actual state instance because migration is different
        // for different state types. For example, ListState needs to deal with
        // individual elements
        StateCreateFactory stateCreateFactory = STATE_CREATE_FACTORIES.get(stateDesc.getType());
        if (stateCreateFactory == null) {
            throw new FlinkRuntimeException(stateNotSupportedMessage(stateDesc));
        }
        State state =
                stateCreateFactory.createState(stateDesc, stateTable, stateTable.keySerializer);
        if (!(state instanceof AbstractHeapState)) {
            throw new FlinkRuntimeException(
                    "State should be an AbstractRocksDBState but is " + state);
        }
        AbstractHeapState<K, N, V> heapState = (AbstractHeapState<K, N, V>) state;
        TtlAwareSerializer<V, ?> currentTtlAwareSerializer =
                (TtlAwareSerializer<V, ?>)
                        TtlAwareSerializer.wrapTtlAwareSerializer(currentSerializer);

        stateTable.transformAll(
                null,
                new StateTransformationFunction<V, V>() {
                    @Override
                    public V apply(V previousState, V value) throws Exception {
                        return heapState.migrateTtlValue(
                                previousState, currentTtlAwareSerializer, ttlTimeProvider);
                    }
                });
    }

    @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);
    }

    @SuppressWarnings("unchecked")
    @Override
    public <N> Stream<K> getKeys(List<String> states, N namespace) {
        final List<StateTable<K, N, ?>> tables =
                states.stream()