private void balance(Map<TopicPartition, ConsumerGenerationPair> prevAssignment,
                             List<TopicPartition> unassignedPartitions) {
            boolean initializing = currentAssignment.get(sortedCurrentSubscriptions.last()).isEmpty();

            // First assign with rack matching and then assign any remaining without rack matching
            List<TopicPartition> partitionsToAssign = unassignedPartitions;
            if (!rackInfo.consumerRacks.isEmpty()) {
                partitionsToAssign = new LinkedList<>(unassignedPartitions);
                maybeAssign(partitionsToAssign, rackInfo, true);
            }
            maybeAssign(partitionsToAssign, null, false);

            // narrow down the reassignment scope to only those partitions that can actually be reassigned
            Set<TopicPartition> fixedPartitions = new HashSet<>();
            for (String topic: topic2AllPotentialConsumers.keySet())
                if (!canTopicParticipateInReassignment(topic)) {
                    for (int i = 0; i < partitionsPerTopic.get(topic).size(); i++) {
                        fixedPartitions.add(new TopicPartition(topic, i));
                    }
                }
            sortedAllPartitions.removeAll(fixedPartitions);
            unassignedPartitions.removeAll(fixedPartitions);

            // narrow down the reassignment scope to only those consumers that are subject to reassignment
            Map<String, List<TopicPartition>> fixedAssignments = new HashMap<>();
            for (String consumer: consumer2AllPotentialTopics.keySet())
                if (!canConsumerParticipateInReassignment(consumer)) {
                    sortedCurrentSubscriptions.remove(consumer);
                    fixedAssignments.put(consumer, currentAssignment.remove(consumer));
                }

            // create a deep copy of the current assignment so we can revert to it if we do not get a more balanced assignment later
            Map<String, List<TopicPartition>> preBalanceAssignment = deepCopy(currentAssignment);
            Map<TopicPartition, String> preBalancePartitionConsumers = new HashMap<>(currentPartitionConsumer);

            // if we don't already need to revoke something due to subscription changes, first try to balance by only moving newly added partitions
            if (!revocationRequired) {
                performReassignments(unassignedPartitions, prevAssignment);
            }

            boolean reassignmentPerformed = performReassignments(sortedAllPartitions, prevAssignment);

            // if we are not preserving existing assignments and we have made changes to the current assignment
            // make sure we are getting a more balanced assignment; otherwise, revert to previous assignment
            if (!initializing && reassignmentPerformed && getBalanceScore(currentAssignment) >= getBalanceScore(preBalanceAssignment)) {
                deepCopy(preBalanceAssignment, currentAssignment);
                currentPartitionConsumer.clear();
                currentPartitionConsumer.putAll(preBalancePartitionConsumers);
            }

            // add the fixed assignments (those that could not change) back
            for (Entry<String, List<TopicPartition>> entry: fixedAssignments.entrySet()) {
                String consumer = entry.getKey();
                currentAssignment.put(consumer, entry.getValue());
                sortedCurrentSubscriptions.add(consumer);
            }

            fixedAssignments.clear();
        }
Balance the current assignment using the data structures created in the assignPartitions(...) method above.
clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignor.java