Source:
private static Map<String, SortedSet<Integer>> getPrevInstanceToTaskAssignments( Iterable<String> liveInstances, ResourceAssignment prevAssignment, Set<Integer> allTaskPartitions, CurrentStateOutput currStateOutput, String jobName, Map<String, Set<Integer>> tasksToDrop) { Map<String, SortedSet<Integer>> result = new HashMap<>(); for (String instance : liveInstances) { result.put(instance, new TreeSet<>()); } for (Partition partition : prevAssignment.getMappedPartitions()) { int pId = TaskUtil.getPartitionId(partition.getPartitionName()); if (allTaskPartitions.contains(pId)) { Map<String, String> replicaMap = prevAssignment.getReplicaMap(partition); for (String instance : replicaMap.keySet()) { SortedSet<Integer> pIdSet = result.get(instance); if (pIdSet != null) { pIdSet.add(pId); } } } } Map<Partition, Map<String, String>> partitions = currStateOutput.getCurrentStateMap(jobName); for (Map.Entry<Partition, Map<String, String>> entry : partitions.entrySet()) { for (Map.Entry<String, String> instanceToCurrState : entry.getValue().entrySet()) { String instance = instanceToCurrState.getKey(); String requestedState = currStateOutput.getRequestedState(jobName, entry.getKey(), instance); TaskPartitionState currState = TaskPartitionState.valueOf(instanceToCurrState.getValue()); int pId = TaskUtil.getPartitionId(entry.getKey().getPartitionName()); if (result.containsKey(instance)) { if (currState == TaskPartitionState.INIT && requestedState != null && requestedState.equals(TaskPartitionState.RUNNING.name()) || currState == TaskPartitionState.RUNNING) { result.get(instance).add(pId); } if (requestedState != null && requestedState.equals(TaskPartitionState.DROPPED.name())) { if (!tasksToDrop.containsKey(instance)) { tasksToDrop.put(instance, new HashSet<>()); } tasksToDrop.get(instance).add(pId); } } } } return result; }

Target:
protected static Map<String, SortedSet<Integer>> getPrevInstanceToTaskAssignments( Iterable<String> liveInstances, ResourceAssignment prevAssignment, Set<Integer> allTaskPartitions, CurrentStateOutput currStateOutput, String jobName, Map<String, Set<Integer>> tasksToDrop) { Map<String, SortedSet<Integer>> result = new HashMap<>(); for (String instance : liveInstances) { result.put(instance, new TreeSet<>()); } for (Partition partition : prevAssignment.getMappedPartitions()) { int pId = TaskUtil.getPartitionId(partition.getPartitionName()); if (allTaskPartitions.contains(pId)) { Map<String, String> replicaMap = prevAssignment.getReplicaMap(partition); for (String instance : replicaMap.keySet()) { SortedSet<Integer> pIdSet = result.get(instance); if (pIdSet != null) { pIdSet.add(pId); } } } } Map<Partition, Map<String, String>> partitions = currStateOutput.getCurrentStateMap(jobName); for (Map.Entry<Partition, Map<String, String>> entry : partitions.entrySet()) { for (Map.Entry<String, String> instanceToCurrState : entry.getValue().entrySet()) { String instance = instanceToCurrState.getKey(); String requestedState = currStateOutput.getRequestedState(jobName, entry.getKey(), instance); TaskPartitionState currState = TaskPartitionState.valueOf(instanceToCurrState.getValue()); int pId = TaskUtil.getPartitionId(entry.getKey().getPartitionName()); if (result.containsKey(instance)) { if (currState == TaskPartitionState.INIT && requestedState != null && requestedState.equals(TaskPartitionState.RUNNING.name()) || currState == TaskPartitionState.RUNNING) { result.get(instance).add(pId); } if (requestedState != null && requestedState.equals(TaskPartitionState.DROPPED.name())) { if (!tasksToDrop.containsKey(instance)) { tasksToDrop.put(instance, new HashSet<>()); } tasksToDrop.get(instance).add(pId); } } } } return result; }