/** * Generate tasks with the assigned topic partitions. * * @param topicGroups group of topics that need to be joined together * @param metadata metadata of the consuming cluster * @return The map from generated task ids to the assigned partitions */ public Map<TaskId, Set<TopicPartition>> partitionGroups(final Map<Integer, Set<String>> topicGroups, final Cluster metadata) { final Map<TaskId, Set<TopicPartition>> groups = new HashMap<>(); for (final Map.Entry<Integer, Set<String>> entry : topicGroups.entrySet()) { final Integer topicGroupId = entry.getKey(); final Set<String> topicGroup = entry.getValue(); final int maxNumPartitions = maxNumPartitions(metadata, topicGroup); for (int partitionId = 0; partitionId < maxNumPartitions; partitionId++) { final Set<TopicPartition> group = new HashSet<>(topicGroup.size()); for (final String topic : topicGroup) { final List<PartitionInfo> partitions = metadata.partitionsForTopic(topic); if (partitionId < partitions.size()) { group.add(new TopicPartition(topic, partitionId)); } } groups.put(new TaskId(topicGroupId, partitionId), Collections.unmodifiableSet(group)); } } return Collections.unmodifiableMap(groups); }