DefaultPartitionGrouper implements PartitionGrouper { public Map<TaskId, Set<TopicPartition>> partitionGroups(Map<Integer, Set<String>> topicGroups, Cluster metadata) { Map<TaskId, Set<TopicPartition>> groups = new HashMap<>(); for (Map.Entry<Integer, Set<String>> entry : topicGroups.entrySet()) { Integer topicGroupId = entry.getKey(); Set<String> topicGroup = entry.getValue(); int maxNumPartitions = maxNumPartitions(metadata, topicGroup); for (int partitionId = 0; partitionId < maxNumPartitions; partitionId++) { Set<TopicPartition> group = new HashSet<>(topicGroup.size()); for (String topic : topicGroup) { if (partitionId < metadata.partitionsForTopic(topic).size()) { group.add(new TopicPartition(topic, partitionId)); } } groups.put(new TaskId(topicGroupId, partitionId), Collections.unmodifiableSet(group)); } } return Collections.unmodifiableMap(groups); }  Map<TaskId, Set<TopicPartition>> partitionGroups(Map<Integer, Set<String>> topicGroups, Cluster metadata);  }
@Test public void testGrouping() { PartitionGrouper grouper = new DefaultPartitionGrouper(); int topicGroupId; Map<TaskId, Set<TopicPartition>> expectedPartitionsForTask; Map<Integer, Set<String>> topicGroups; topicGroups = new HashMap<>(); topicGroups.put(0, mkSet("topic1")); topicGroups.put(1, mkSet("topic2")); expectedPartitionsForTask = new HashMap<>(); topicGroupId = 0; expectedPartitionsForTask.put(new TaskId(topicGroupId, 0), mkSet(new TopicPartition("topic1", 0))); expectedPartitionsForTask.put(new TaskId(topicGroupId, 1), mkSet(new TopicPartition("topic1", 1))); expectedPartitionsForTask.put(new TaskId(topicGroupId, 2), mkSet(new TopicPartition("topic1", 2))); topicGroupId++; expectedPartitionsForTask.put(new TaskId(topicGroupId, 0), mkSet(new TopicPartition("topic2", 0))); expectedPartitionsForTask.put(new TaskId(topicGroupId, 1), mkSet(new TopicPartition("topic2", 1))); assertEquals(expectedPartitionsForTask, grouper.partitionGroups(topicGroups, metadata)); topicGroups = new HashMap<>(); topicGroups.put(0, mkSet("topic1", "topic2")); expectedPartitionsForTask = new HashMap<>(); topicGroupId = 0; expectedPartitionsForTask.put(new TaskId(topicGroupId, 0), mkSet(new TopicPartition("topic1", 0), new TopicPartition("topic2", 0))); expectedPartitionsForTask.put(new TaskId(topicGroupId, 1), mkSet(new TopicPartition("topic1", 1), new TopicPartition("topic2", 1))); expectedPartitionsForTask.put(new TaskId(topicGroupId, 2), mkSet(new TopicPartition("topic1", 2))); assertEquals(expectedPartitionsForTask, grouper.partitionGroups(topicGroups, metadata)); }