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 testStateWithFullyFinishedOperators() throws JobException, JobExecutionException {
        List<OperatorID> operatorIds = buildOperatorIds(2);
        Map<OperatorID, OperatorState> states =
                buildOperatorStates(Collections.singletonList(operatorIds.get(1)), 3);

        // Create an operator state marked as finished
        OperatorState operatorState = new FullyFinishedOperatorState(operatorIds.get(0), 3, 256);
        states.put(operatorIds.get(0), operatorState);

        Map<OperatorID, ExecutionJobVertex> vertices =
                buildVertices(operatorIds, 2, RANGE, ROUND_ROBIN);
        new StateAssignmentOperation(0, new HashSet<>(vertices.values()), states, false)
                .assignStates();

        // Check the job vertex with only finished operator.
        ExecutionJobVertex jobVertexWithFinishedOperator = vertices.get(operatorIds.get(0));
        for (ExecutionVertex task : jobVertexWithFinishedOperator.getTaskVertices()) {
            JobManagerTaskRestore taskRestore = task.getCurrentExecutionAttempt().getTaskRestore();
            assertThat(taskRestore.getTaskStateSnapshot().isTaskDeployedAsFinished()).isTrue();
        }

        // Check the job vertex without finished operator.
        ExecutionJobVertex jobVertexWithoutFinishedOperator = vertices.get(operatorIds.get(1));
        for (ExecutionVertex task : jobVertexWithoutFinishedOperator.getTaskVertices()) {
            JobManagerTaskRestore taskRestore = task.getCurrentExecutionAttempt().getTaskRestore();
            assertThat(taskRestore.getTaskStateSnapshot().isTaskDeployedAsFinished()).isFalse();
        }
    }