public Map<TopicPartition, StopReplicaPartitionState> partitionStates() {
        Map<TopicPartition, StopReplicaPartitionState> partitionStates = new HashMap<>();

        if (version() < 1) {
            for (StopReplicaPartitionV0 partition : data.ungroupedPartitions()) {
                partitionStates.put(
                    new TopicPartition(partition.topicName(), partition.partitionIndex()),
                    new StopReplicaPartitionState()
                        .setPartitionIndex(partition.partitionIndex())
                        .setDeletePartition(data.deletePartitions()));
            }
        } else if (version() < 3) {
            for (StopReplicaTopicV1 topic : data.topics()) {
                for (Integer partitionIndex : topic.partitionIndexes()) {
                    partitionStates.put(
                        new TopicPartition(topic.name(), partitionIndex),
                        new StopReplicaPartitionState()
                            .setPartitionIndex(partitionIndex)
                            .setDeletePartition(data.deletePartitions()));
                }
            }
        } else {
            for (StopReplicaTopicState topicState : data.topicStates()) {
                for (StopReplicaPartitionState partitionState: topicState.partitionStates()) {
                    partitionStates.put(
                        new TopicPartition(topicState.topicName(), partitionState.partitionIndex()),
                        partitionState);
                }
            }
        }

        return partitionStates;
    }
Note that this method has allocation overhead per iterated element, so callers should copy the result into
another collection if they need to iterate more than once.

Implementation note: we should strive to avoid allocation overhead per element, see
`UpdateMetadataRequest.partitionStates()` for the preferred approach. That's not possible in this case and
StopReplicaRequest should be relatively rare in comparison to other request types.
clients/src/main/java/org/apache/kafka/common/requests/StopReplicaRequest.java