GroupByPartition implements SystemStreamPartitionGrouper { @Override public Map<TaskName, Set<SystemStreamPartition>> group(Set<SystemStreamPartition> ssps) { Map<TaskName, Set<SystemStreamPartition>> groupedMap = new HashMap<>(); for (SystemStreamPartition ssp : ssps) { if (broadcastStreams.contains(ssp)) { continue; } TaskName taskName = new TaskName(String.format("Partition %d", ssp.getPartition().getPartitionId())); if (!groupedMap.containsKey(taskName)) { groupedMap.put(taskName, new HashSet<>()); } groupedMap.get(taskName).add(ssp); } if (!broadcastStreams.isEmpty()) { for (Set<SystemStreamPartition> value : groupedMap.values()) { value.addAll(broadcastStreams); } } return groupedMap; } GroupByPartition(Config config); @Override Map<TaskName, Set<SystemStreamPartition>> group(Set<SystemStreamPartition> ssps);  }
@Test public void testLocalStreamsGroupedCorrectly() { GroupByPartition grouper = new GroupByPartition(new MapConfig()); Map<TaskName, Set<SystemStreamPartition>> emptyResult = grouper.group(new HashSet<>()); assertTrue(emptyResult.isEmpty()); Map<TaskName, Set<SystemStreamPartition>> result = grouper.group(ImmutableSet.of(aa0, aa1, aa2, ab1, ab2, ac0)); Map<TaskName, Set<SystemStreamPartition>> expectedResult = ImmutableMap.<TaskName, Set<SystemStreamPartition>>builder() .put(new TaskName("Partition 0"), ImmutableSet.of(aa0, ac0)) .put(new TaskName("Partition 1"), ImmutableSet.of(aa1, ab1)) .put(new TaskName("Partition 2"), ImmutableSet.of(aa2, ab2)) .build(); assertEquals(expectedResult, result); }
@Test public void testBroadcastStreamsGroupedCorrectly() { Config config = new MapConfig(ImmutableMap.of("task.broadcast.inputs", "SystemA.StreamA#0, SystemA.StreamB#1")); GroupByPartition grouper = new GroupByPartition(config); Map<TaskName, Set<SystemStreamPartition>> result = grouper.group(ImmutableSet.of(aa0, aa1, aa2, ab1, ab2, ac0)); Map<TaskName, Set<SystemStreamPartition>> expectedResult = ImmutableMap.<TaskName, Set<SystemStreamPartition>>builder() .put(new TaskName("Partition 0"), ImmutableSet.of(aa0, ac0, ab1)) .put(new TaskName("Partition 1"), ImmutableSet.of(aa1, aa0, ab1)) .put(new TaskName("Partition 2"), ImmutableSet.of(aa2, aa0, ab2, ab1)) .build(); assertEquals(expectedResult, result); }
@Test public void testNoTaskOnlyContainsBroadcastStreams() { Config config = new MapConfig(ImmutableMap.of("task.broadcast.inputs", "SystemA.StreamA#0, SystemA.StreamB#1")); GroupByPartition grouper = new GroupByPartition(config); Map<TaskName, Set<SystemStreamPartition>> result = grouper.group(ImmutableSet.of(aa0, ab1, ab2)); Map<TaskName, Set<SystemStreamPartition>> expectedResult = ImmutableMap.<TaskName, Set<SystemStreamPartition>>builder() .put(new TaskName("Partition 2"), ImmutableSet.of(aa0, ab1, ab2)).build(); assertEquals(expectedResult, result); }