Repository: apache/kafka
File: tools/src/main/java/org/apache/kafka/tools/consumer/group/ConsumerGroupCommand.java (Java)

Code before:
        private Optional<Long> getLag(Optional<Long> offset, Optional<Long> logEndOffset) {
            return offset.filter(o -> o != -1).flatMap(offset0 -> logEndOffset.map(end -> end - offset0));
        }

        private Collection<PartitionAssignmentState> describePartitions(
            String group,
            Optional<Node> coordinator,
            List<TopicPartition> topicPartitions,
            Map<TopicPartition, OffsetAndMetadata> committedOffsets,
            Optional<String> consumerIdOpt,
            Optional<String> hostOpt,
            Optional<String> clientIdOpt
        ) {
            BiFunction<TopicPartition, Optional<Long>, PartitionAssignmentState> getDescribePartitionResult = (topicPartition, logEndOffsetOpt) -> {
                // The admin client returns `null` as a value to indicate that there is not committed offset for a partition.
                Optional<Long> offset = Optional.ofNullable(committedOffsets.get(topicPartition)).map(OffsetAndMetadata::offset);
                Optional<Integer> leaderEpoch = Optional.ofNullable(committedOffsets.get(topicPartition)).flatMap(OffsetAndMetadata::leaderEpoch);
                return new PartitionAssignmentState(group, coordinator, Optional.of(topicPartition.topic()),
                    Optional.of(topicPartition.partition()), offset, getLag(offset, logEndOffsetOpt),
                    consumerIdOpt, hostOpt, clientIdOpt, logEndOffsetOpt, leaderEpoch);
            };

            Set<TopicPartition> nonLeaderTopicPartitions = filterNoneLeaderPartitions(topicPartitions);

            // prepare data for partitions with leaders
            topicPartitions.removeAll(nonLeaderTopicPartitions);
            List<PartitionAssignmentState> existLeaderAssignments = offsetsUtils.getLogEndOffsets(topicPartitions).entrySet().stream().map(logEndOffsetResult -> {
                if (logEndOffsetResult.getValue() instanceof OffsetsUtils.LogOffset)
                    return getDescribePartitionResult.apply(
                        logEndOffsetResult.getKey(),
                        Optional.of(((OffsetsUtils.LogOffset) logEndOffsetResult.getValue()).value())
                    );
                else if (logEndOffsetResult.getValue() instanceof OffsetsUtils.Unknown)
                    return getDescribePartitionResult.apply(logEndOffsetResult.getKey(), Optional.empty());
                else if (logEndOffsetResult.getValue() instanceof OffsetsUtils.Ignore)
                    return null;

                throw new IllegalStateException("Unknown LogOffset subclass: " + logEndOffsetResult.getValue());
            }).toList();

            // prepare data for partitions without leaders
            List<PartitionAssignmentState> noneLeaderAssignments = nonLeaderTopicPartitions.stream()
                    .map(tp -> getDescribePartitionResult.apply(tp, Optional.empty())).toList();

            // concat the data and then sort them
            return Stream.concat(existLeaderAssignments.stream(), noneLeaderAssignments.stream())
                    .sorted(Comparator.<PartitionAssignmentState, String>comparing(
                            state -> state.topic.orElse(""), String::compareTo)
                            .thenComparingInt(state -> state.partition.orElse(-1)))
                    .collect(Collectors.toList());
        }

Review comment (suggestion):
we don't need to modify the input `topicPartitions` if `filterNoneLeaderPartitions` returns the partitions having leaders.

Code after:
            String group,
            Optional<Node> coordinator,
            Collection<TopicPartition> topicPartitions,
            Map<TopicPartition, OffsetAndMetadata> committedOffsets,
            Optional<String> consumerIdOpt,
            Optional<String> hostOpt,
            Optional<String> clientIdOpt
        ) {
            BiFunction<TopicPartition, Optional<Long>, PartitionAssignmentState> getDescribePartitionResult = (topicPartition, logEndOffsetOpt) -> {
                // The admin client returns `null` as a value to indicate that there is not committed offset for a partition.
                Optional<Long> offset = Optional.ofNullable(committedOffsets.get(topicPartition)).map(OffsetAndMetadata::offset);
                Optional<Integer> leaderEpoch = Optional.ofNullable(committedOffsets.get(topicPartition)).flatMap(OffsetAndMetadata::leaderEpoch);
                return new PartitionAssignmentState(group, coordinator, Optional.of(topicPartition.topic()),
                    Optional.of(topicPartition.partition()), offset, getLag(offset, logEndOffsetOpt),
                    consumerIdOpt, hostOpt, clientIdOpt, logEndOffsetOpt, leaderEpoch);
            };

            List<TopicPartition> topicPartitionsWithoutLeader = filterNoneLeaderPartitions(topicPartitions);
            List<TopicPartition> topicPartitionsWithLeader = topicPartitions.stream().filter(tp -> !topicPartitionsWithoutLeader.contains(tp)).toList();

            // prepare data for partitions with leaders
            List<PartitionAssignmentState> existLeaderAssignments = offsetsUtils.getLogEndOffsets(topicPartitionsWithLeader).entrySet().stream().map(logEndOffsetResult -> {
                if (logEndOffsetResult.getValue() instanceof OffsetsUtils.LogOffset)
                    return getDescribePartitionResult.apply(
                        logEndOffsetResult.getKey(),
                        Optional.of(((OffsetsUtils.LogOffset) logEndOffsetResult.getValue()).value())
                    );
                else if (logEndOffsetResult.getValue() instanceof OffsetsUtils.Unknown)
                    return getDescribePartitionResult.apply(logEndOffsetResult.getKey(), Optional.empty());
                else if (logEndOffsetResult.getValue() instanceof OffsetsUtils.Ignore)
                    return null;

                throw new IllegalStateException("Unknown LogOffset subclass: " + logEndOffsetResult.getValue());
            }).toList();

            // prepare data for partitions without leaders
            List<PartitionAssignmentState> noneLeaderAssignments = topicPartitionsWithoutLeader.stream()
                    .map(tp -> getDescribePartitionResult.apply(tp, Optional.empty())).toList();

            // concat the data and then sort them
            return Stream.concat(existLeaderAssignments.stream(), noneLeaderAssignments.stream())
                    .sorted(Comparator.<PartitionAssignmentState, String>comparing(
                            state -> state.topic.orElse(""), String::compareTo)
                            .thenComparingInt(state -> state.partition.orElse(-1)))
                    .collect(Collectors.toList());
        }

        private List<TopicPartition> filterNoneLeaderPartitions(Collection<TopicPartition> topicPartitions) {
            // collect all topics
            Set<String> topics = topicPartitions.stream().map(TopicPartition::topic).collect(Collectors.toSet());