KafkaUtils { public static List<Partition> calculatePartitionsForTask(GlobalPartitionInformation partitionInformation, int totalTasks, int taskIndex) { Preconditions.checkArgument(taskIndex < totalTasks, "task index must be less that total tasks"); List<Partition> partitions = partitionInformation.getOrderedPartitions(); int numPartitions = partitions.size(); if (numPartitions < totalTasks) { LOG.warn("there are more tasks than partitions (tasks: " + totalTasks + "; partitions: " + numPartitions + "), some tasks will be idle"); } List<Partition> taskPartitions = new ArrayList<Partition>(); for (int i = taskIndex; i < numPartitions; i += totalTasks) { Partition taskPartition = partitions.get(i); taskPartitions.add(taskPartition); } logPartitionMapping(totalTasks, taskIndex, taskPartitions); return taskPartitions; }  static IBrokerReader makeBrokerReader(Map stormConf, KafkaConfig conf); static long getOffset(SimpleConsumer consumer, String topic, int partition, KafkaConfig config); static long getOffset(SimpleConsumer consumer, String topic, int partition, long startOffsetTime); static ByteBufferMessageSet fetchMessages(KafkaConfig config, SimpleConsumer consumer, Partition partition, long offset); static Iterable<List<Object>> generateTuples(KafkaConfig kafkaConfig, Message msg); static List<Partition> calculatePartitionsForTask(GlobalPartitionInformation partitionInformation, int totalTasks, int taskIndex); static String taskId(int taskIndex, int totalTasks); static final Logger LOG; }
@Test public void moreTasksThanPartitions() { GlobalPartitionInformation globalPartitionInformation = TestUtils.buildPartitionInfo(1); int numTasks = 2; assertEquals(1, KafkaUtils.calculatePartitionsForTask(globalPartitionInformation, numTasks, 0).size()); assertEquals(0, KafkaUtils.calculatePartitionsForTask(globalPartitionInformation, numTasks, 1).size()); }
@Test (expected = IllegalArgumentException.class ) public void assignInvalidTask() { GlobalPartitionInformation globalPartitionInformation = new GlobalPartitionInformation(); KafkaUtils.calculatePartitionsForTask(globalPartitionInformation, 1, 1); }