Repository: apache/kafka
File: tools/src/main/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommand.java (Java)

Code before:
    }

    /**
     * The entry point for the --generate command.
     *
     * @param adminClient           The AdminClient to use.
     * @param reassignmentJson      The JSON string to use for the topics to reassign.
     * @param brokerListString      The comma-separated string of broker IDs to use.
     * @param enableRackAwareness   True if rack-awareness should be enabled.
     *
     * @return                      A tuple containing the proposed assignment and the
     *                              current assignment.
     */
    public static Entry<Map<TopicPartition, List<Integer>>, Map<TopicPartition, List<Integer>>> generateAssignment(Admin adminClient,
                                                                                                                   String reassignmentJson,
                                                                                                                   String brokerListString,
                                                                                                                   boolean enableRackAwareness
    ) throws ExecutionException, InterruptedException, JsonProcessingException {
        Entry<List<Integer>, List<String>> t0 = parseGenerateAssignmentArgs(reassignmentJson, brokerListString);

        List<Integer> brokersToReassign = t0.getKey();
        List<String> topicsToReassign = t0.getValue();

        Map<TopicPartition, List<Node>> currentAssignments = getReplicaAssignmentForTopics(adminClient, topicsToReassign);
        Map<TopicPartition, List<Integer>> currentParts = toReplicaIds(currentAssignments);
        Map<TopicPartition, List<Integer>> currentActiveParts = new HashMap<>();
        Map<TopicPartition, List<Integer>> currentInactiveParts = new HashMap<>();
        splitReplicaIds(currentAssignments, currentActiveParts, currentInactiveParts);
        Map<TopicPartitionReplica, String> currentReplicaLogDirs = getReplicaToLogDir(adminClient, currentActiveParts, currentInactiveParts);
        List<UsableBroker> usableBrokers = getBrokerMetadata(adminClient, brokersToReassign, enableRackAwareness);
        Map<TopicPartition, List<Integer>> proposedAssignments = calculateAssignment(currentParts, usableBrokers);
        System.out.printf("Current partition replica assignment%n%s%n%n",
            formatAsReassignmentJson(currentParts, currentReplicaLogDirs));
        System.out.printf("Proposed partition reassignment configuration%n%s%n",
            formatAsReassignmentJson(proposedAssignments, Map.of()));
        return Map.entry(proposedAssignments, currentParts);
    }

    /**
     * Calculate the new partition assignments to suggest in --generate.
     *
     * @param currentAssignment  The current partition assignments.
     * @param usableBrokers      The rack information for each broker.
     *
     * @return                   A map from partitions to the proposed assignments for each.
     */
    private static Map<TopicPartition, List<Integer>> calculateAssignment(Map<TopicPartition, List<Integer>> currentAssignment,
                                                                          List<UsableBroker> usableBrokers) {
        Map<String, List<Entry<TopicPartition, List<Integer>>>> groupedByTopic = new HashMap<>();
        for (Entry<TopicPartition, List<Integer>> e : currentAssignment.entrySet())
            groupedByTopic.computeIfAbsent(e.getKey().topic(), k -> new ArrayList<>()).add(e);

Review comment (refactor):
The style feels a bit awkward, since we have to create an extra collection before we use it. Maybe we can add a helper method instead. For example:
```java
    static List<TopicPartitionReplica> available(Map<TopicPartition, List<Node>> current) {
        return current.entrySet().stream().flatMap(entry -> entry.getValue()
                        .stream().filter(n -> !n.isEmpty()).map(n -> new TopicPartitionReplica(entry.getKey().topic(), entry.getKey().partition(), n.id())))
                .toList();
    }

    static Map<TopicPartitionReplica, String> getReplicaToLogDir(
            Admin adminClient, Map<TopicPartition, List<Node>> current) throws ExecutionException, InterruptedException {
        return adminClient.describeReplicaLogDirs(available(current)).all().get().entrySet()
                .stream()
                .filter(e -> e.getValue().getCurrentReplicaLogDir() != null)
                .collect(Collectors.toMap(Entry::getKey, e -> e.getValue().getCurrentReplicaLogDir()));
    }
```

Code after:
    }

    /**
     * The entry point for the --generate command.
     *
     * @param adminClient           The AdminClient to use.
     * @param reassignmentJson      The JSON string to use for the topics to reassign.
     * @param brokerListString      The comma-separated string of broker IDs to use.
     * @param enableRackAwareness   True if rack-awareness should be enabled.
     *
     * @return                      A tuple containing the proposed assignment and the
     *                              current assignment.
     */
    public static Entry<Map<TopicPartition, List<Integer>>, Map<TopicPartition, List<Integer>>> generateAssignment(Admin adminClient,
                                                                                                                   String reassignmentJson,
                                                                                                                   String brokerListString,
                                                                                                                   boolean enableRackAwareness
    ) throws ExecutionException, InterruptedException, JsonProcessingException {
        Entry<List<Integer>, List<String>> t0 = parseGenerateAssignmentArgs(reassignmentJson, brokerListString);

        List<Integer> brokersToReassign = t0.getKey();
        List<String> topicsToReassign = t0.getValue();

        Map<TopicPartition, List<Node>> currentAssignments = getReplicaAssignmentForTopics(adminClient, topicsToReassign);
        Map<TopicPartitionReplica, String> currentReplicaLogDirs = getReplicaToLogDir(adminClient, currentAssignments);
        List<UsableBroker> usableBrokers = getBrokerMetadata(adminClient, brokersToReassign, enableRackAwareness);
        Map<TopicPartition, List<Integer>> currentParts = toReplicaIds(currentAssignments);
        Map<TopicPartition, List<Integer>> proposedAssignments = calculateAssignment(currentParts, usableBrokers);
        System.out.printf("Current partition replica assignment%n%s%n%n",
            formatAsReassignmentJson(currentParts, currentReplicaLogDirs));
        System.out.printf("Proposed partition reassignment configuration%n%s%n",
            formatAsReassignmentJson(proposedAssignments, Map.of()));
        return Map.entry(proposedAssignments, currentParts);
    }

    /**
     * Calculate the new partition assignments to suggest in --generate.
     *
     * @param currentAssignment  The current partition assignments.
     * @param usableBrokers      The rack information for each broker.
     *
     * @return                   A map from partitions to the proposed assignments for each.
     */
    private static Map<TopicPartition, List<Integer>> calculateAssignment(Map<TopicPartition, List<Integer>> currentAssignment,
                                                                          List<UsableBroker> usableBrokers) {
        Map<String, List<Entry<TopicPartition, List<Integer>>>> groupedByTopic = new HashMap<>();
        for (Entry<TopicPartition, List<Integer>> e : currentAssignment.entrySet())
            groupedByTopic.computeIfAbsent(e.getKey().topic(), k -> new ArrayList<>()).add(e);
        Map<TopicPartition, List<Integer>> proposedAssignments = new HashMap<>();
        groupedByTopic.forEach((topic, assignment) -> {
            List<Integer> replicas = assignment.get(0).getValue();