Source:
<code>public void assignStates() { Map<OperatorID, OperatorState> localOperators = new HashMap<>(operatorStates); checkStateMappingCompleteness(allowNonRestoredState, operatorStates, tasks); for (ExecutionJobVertex executionJobVertex : this.tasks) { List<OperatorIDPair> operatorIDPairs = executionJobVertex.getOperatorIDs(); Map<OperatorID, OperatorState> operatorStates = new HashMap<>(operatorIDPairs.size()); for (OperatorIDPair operatorIDPair : operatorIDPairs) { OperatorID operatorID = operatorIDPair.getUserDefinedOperatorID().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); vertexAssignments.put(executionJobVertex, stateAssignment); for (final IntermediateResult producedDataSet : executionJobVertex.getInputs()) { <START> consumerAssignment.put(producedDataSet, stateAssignment); <END> } } for (TaskStateAssignment stateAssignment : vertexAssignments.values()) { if (stateAssignment.hasState) { assignAttemptState(stateAssignment); } } for (TaskStateAssignment stateAssignment : vertexAssignments.values()) { if (stateAssignment.hasState) { assignTaskStateToExecutionJobVertices(stateAssignment); } } }</code><technical_language>1. I if this put override existing assignment (e.g. a UnionGate). 2. IntermediateResult a key override equals/hashCode - is intentional</technical_language>

Target:
public void assignStates() { Map<OperatorID, OperatorState> localOperators = new HashMap<>(operatorStates); checkStateMappingCompleteness(allowNonRestoredState, operatorStates, tasks); for (ExecutionJobVertex executionJobVertex : this.tasks) { List<OperatorIDPair> operatorIDPairs = executionJobVertex.getOperatorIDs(); Map<OperatorID, OperatorState> operatorStates = new HashMap<>(operatorIDPairs.size()); for (OperatorIDPair operatorIDPair : operatorIDPairs) { OperatorID operatorID = operatorIDPair.getUserDefinedOperatorID().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); vertexAssignments.put(executionJobVertex, stateAssignment); for (final IntermediateResult producedDataSet : executionJobVertex.getInputs()) { consumerAssignment.put(producedDataSet.getId(), stateAssignment); } } for (TaskStateAssignment stateAssignment : vertexAssignments.values()) { if (stateAssignment.hasState) { assignAttemptState(stateAssignment); } } for (TaskStateAssignment stateAssignment : vertexAssignments.values()) { if (stateAssignment.hasState) { assignTaskStateToExecutionJobVertices(stateAssignment); } } }