private void assignWithRackMatching(Collection<TopicAssignmentState> assignmentStates,
                                        Map<String, List<TopicPartition>> assignment) {

        assignmentStates.stream().collect(Collectors.groupingBy(t -> t.consumers)).forEach((consumers, states) ->
            states.stream().collect(Collectors.groupingBy(t -> t.partitionRacks.size())).forEach((numPartitions, coPartitionedStates) -> {
                if (coPartitionedStates.size() > 1)
                    assignCoPartitionedWithRackMatching(consumers, numPartitions, coPartitionedStates, assignment);
                else {
                    TopicAssignmentState state = coPartitionedStates.get(0);
                    if (state.needsRackAwareAssignment)
                        assignRanges(state, state::racksMatch, assignment);
                }
            })
        );
    }
Performs range assignment of the specified partitions for the consumers with the provided subscriptions.
If rack-awareness is enabled for one or more consumers, we perform rack-aware assignment first to assign
the subset of partitions that can be aligned on racks, while retaining the same co-partitioning and
per-topic balancing guarantees as non-rack-aware range assignment. The remaining partitions are assigned
using standard non-rack-aware range assignment logic, which may result in mis-aligned racks.
clients/src/main/java/org/apache/kafka/clients/consumer/RangeAssignor.java