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

Code before:
            try {
                future.get();
                results.add(replica);
            } catch (ExecutionException t) {
                // Ignore ReplicaNotAvailableException.  It is OK if the replica is not
                // available at this moment.
                if (t.getCause() instanceof ReplicaNotAvailableException)
                    continue;
                throw new AdminCommandFailedException("Failed to alter dir for " + replica, t);
            }
        }
        return results;
    }

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

    private static List<TopicPartitionReplica> getTopicPartitionReplicas(Map<TopicPartition, List<Integer>> topicPartitionToReplicas) {
        return topicPartitionToReplicas.entrySet().stream()
                .flatMap(entry -> {
                    TopicPartition tp = entry.getKey();
                    List<Integer> brokerIds = entry.getValue();
                    return brokerIds.stream()
                            .map(brokerId -> new TopicPartitionReplica(tp.topic(), tp.partition(), brokerId));
                })
                .toList();
    }
}


Review comment (refactor):
It seems `getTopicPartitionReplicas` and `getReplicaToLogDir` can be merged into a single helper.
```java
    static Map<TopicPartitionReplica, String> getReplicaToLogDir(Admin adminClient,
                                                                 Map<TopicPartition, List<Integer>> topicPartitionToReplicas) throws InterruptedException, ExecutionException {
        var tprs = topicPartitionToReplicas.entrySet().stream()
                .flatMap(entry -> entry.getValue().stream()
                        .map(id -> new TopicPartitionReplica(entry.getKey().topic(), entry.getKey().partition(), id)))
                .collect(Collectors.toUnmodifiableSet());

        return adminClient.describeReplicaLogDirs(tprs).all().get()
                .entrySet()
                .stream()
                .filter(entry -> entry.getValue().getCurrentReplicaLogDir() != null)
                .collect(Collectors.toMap(
                        Entry::getKey,
                        entry -> entry.getValue().getCurrentReplicaLogDir()
                ));
    }

```

Code after:
                results.add(replica);
            } catch (ExecutionException t) {
                // Ignore ReplicaNotAvailableException.  It is OK if the replica is not
                // available at this moment.
                if (t.getCause() instanceof ReplicaNotAvailableException)
                    continue;
                throw new AdminCommandFailedException("Failed to alter dir for " + replica, t);
            }
        }
        return results;
    }

    static Map<TopicPartitionReplica, String> getReplicaToLogDir(
        Admin adminClient,
        Map<TopicPartition, List<Integer>> topicPartitionToReplicas
    ) throws InterruptedException, ExecutionException {
        var replicaLogDirs = topicPartitionToReplicas
                .entrySet()
                .stream()
                .flatMap(entry -> entry.getValue()
                    .stream()
                    .map(id -> new TopicPartitionReplica(entry.getKey().topic(), entry.getKey().partition(), id)))
                .collect(Collectors.toUnmodifiableSet());

        return adminClient.describeReplicaLogDirs(replicaLogDirs).all().get()
                .entrySet()
                .stream()
                .filter(entry -> entry.getValue().getCurrentReplicaLogDir() != null)
                .collect(Collectors.toMap(
                    Entry::getKey,
                    entry -> entry.getValue().getCurrentReplicaLogDir()
                ));
    }
}