public void assignStates() {
        checkStateMappingCompleteness(allowNonRestoredState, operatorStates, tasks);

        Map<OperatorID, OperatorState> localOperators = new HashMap<>(operatorStates);

        // find the states of all operators belonging to this task and compute additional
        // information in first pass
        for (ExecutionJobVertex executionJobVertex : tasks) {
            List<OperatorIDPair> operatorIDPairs = executionJobVertex.getOperatorIDs();
            Map<OperatorID, OperatorState> operatorStates =
                    CollectionUtil.newHashMapWithExpectedSize(operatorIDPairs.size());
            for (OperatorIDPair operatorIDPair : operatorIDPairs) {
                OperatorID operatorID =
                        operatorIDPair
                                .getUserDefinedOperatorID()
                                .filter(localOperators::containsKey)
                                .orElse(operatorIDPair.getGeneratedOperatorID());

                OperatorState operatorState = localOperators.remove(operatorID);
                if (operatorState == null) {
                    operatorState =
                            new OperatorState(
                                    operatorID,
                                    executionJobVertex.getParallelism(),
                                    executionJobVertex.getMaxParallelism());
                }
                operatorStates.put(operatorIDPair.getGeneratedOperatorID(), operatorState);
            }

            final TaskStateAssignment stateAssignment =
                    new TaskStateAssignment(
                            executionJobVertex,
                            operatorStates,
                            consumerAssignment,
                            vertexAssignments);
            vertexAssignments.put(executionJobVertex, stateAssignment);
            for (final IntermediateResult producedDataSet : executionJobVertex.getInputs()) {
                consumerAssignment.put(producedDataSet.getId(), stateAssignment);
            }
        }

        // repartition state
        for (TaskStateAssignment stateAssignment : vertexAssignments.values()) {
            if (stateAssignment.hasNonFinishedState
                    // FLINK-31963: We need to run repartitioning for stateless operators that have
                    // upstream output or downstream input states.
                    || stateAssignment.hasUpstreamOutputStates()
                    || stateAssignment.hasDownstreamInputStates()) {
                assignAttemptState(stateAssignment);
            }
        }

        // actually assign the state
        for (TaskStateAssignment stateAssignment : vertexAssignments.values()) {
            // If upstream has output states or downstream has input states, even the empty task
            // state should be assigned for the current task in order to notify this task that the
            // old states will send to it which likely should be filtered.
            if (stateAssignment.hasNonFinishedState
                    || stateAssignment.isFullyFinished
                    || stateAssignment.hasUpstreamOutputStates()
                    || stateAssignment.hasDownstreamInputStates()) {
                assignTaskStateToExecutionJobVertices(stateAssignment);
            }
        }
    }
@Test
    void testChannelStateAssignmentUpscaling() throws JobException, JobExecutionException {
        List<OperatorID> operatorIds = buildOperatorIds(2);
        Map<OperatorID, OperatorState> states = buildOperatorStates(operatorIds, 2);

        Map<OperatorID, ExecutionJobVertex> vertices =
                buildVertices(operatorIds, 3, RANGE, ROUND_ROBIN);

        new StateAssignmentOperation(0, new HashSet<>(vertices.values()), states, false)
                .assignStates();

        for (OperatorID operatorId : operatorIds) {
            // input is range partitioned, so there is an overlap
            assertState(
                    vertices, operatorId, states, 0, OperatorSubtaskState::getInputChannelState, 0);
            assertState(
                    vertices,
                    operatorId,
                    states,
                    1,
                    OperatorSubtaskState::getInputChannelState,
                    0,
                    1);
            assertState(
                    vertices, operatorId, states, 2, OperatorSubtaskState::getInputChannelState, 1);
            // output is round robin redistributed
            assertState(
                    vertices,
                    operatorId,
                    states,
                    0,
                    OperatorSubtaskState::getResultSubpartitionState,
                    0);
            assertState(
                    vertices,
                    operatorId,
                    states,
                    1,
                    OperatorSubtaskState::getResultSubpartitionState,
                    1);
            assertState(
                    vertices,
                    operatorId,
                    states,
                    2,
                    OperatorSubtaskState::getResultSubpartitionState);
        }

        assertThat(
                        getAssignedState(vertices.get(operatorIds.get(0)), operatorIds.get(0), 0)
                                .getOutputRescalingDescriptor())
                .isEqualTo(
                        rescalingDescriptor(to(0), array(mappings(to(0), to(0, 1), to(1))), set()));
        assertThat(
                        getAssignedState(vertices.get(operatorIds.get(0)), operatorIds.get(0), 1)
                                .getOutputRescalingDescriptor())
                .isEqualTo(
                        rescalingDescriptor(to(1), array(mappings(to(0), to(0, 1), to(1))), set()));
        // unmapped subtask index, so nothing to do
        assertThat(
                        getAssignedState(vertices.get(operatorIds.get(0)), operatorIds.get(0), 2)
                                .getOutputRescalingDescriptor())
                .isEqualTo(InflightDataRescalingDescriptor.NO_RESCALE);

        assertThat(
                        getAssignedState(vertices.get(operatorIds.get(1)), operatorIds.get(1), 0)
                                .getInputRescalingDescriptor())
                .isEqualTo(
                        rescalingDescriptor(to(0), array(mappings(to(0), to(1), to())), set(0, 1)));
        assertThat(
                        getAssignedState(vertices.get(operatorIds.get(1)), operatorIds.get(1), 1)
                                .getInputRescalingDescriptor())
                .isEqualTo(
                        rescalingDescriptor(
                                to(0, 1), array(mappings(to(0), to(1), to())), set(0, 1)));
        assertThat(
                        getAssignedState(vertices.get(operatorIds.get(1)), operatorIds.get(1), 2)
                                .getInputRescalingDescriptor())
                .isEqualTo(
                        rescalingDescriptor(to(1), array(mappings(to(0), to(1), to())), set(0, 1)));
    }