GroupPartitioner { public List<Integer> getPartitionIds(final Object[] partialKey) { if (null == partialKey) { throw new IllegalArgumentException("getPartitionIds cannot be called with null partialKey"); } if (splitPoints.isEmpty()) { return Collections.singletonList(0); } final int partitionKeyLength = splitPoints.get(0).getLength(); if (partialKey.length == partitionKeyLength) { return getPartitionIdsSameLengthKey(partialKey); } if (partialKey.length < partitionKeyLength) { return getPartitionIdsShorterKey(partialKey); } throw new IllegalArgumentException("getPartitionIds cannot be called with a partialKey that is longer than the split points"); } GroupPartitioner(final String group, final List<PartitionKey> splitPoints); String getGroup(); int getPartitionId(final PartitionKey partitionKey); List<Integer> getPartitionIds(final Object[] partialKey); PartitionKey getIthPartitionKey(final int i); List<PartitionKey> getSplitPoints(); List<Partition> getPartitions(); @Override String toString(); @Override boolean equals(final Object obj); @Override int hashCode();  }
@Test public void testGetPartitionIds() { final PartitionKey key1 = new PartitionKey(new Object[]{1L, 4L}); final PartitionKey key2 = new PartitionKey(new Object[]{1L, 7L}); final PartitionKey key3 = new PartitionKey(new Object[]{1L, 9L}); final PartitionKey key4 = new PartitionKey(new Object[]{5L, 4L}); final List<PartitionKey> splitPoints = new ArrayList<>(Arrays.asList(key1, key2, key3, key4)); final GroupPartitioner partitioner = new GroupPartitioner("GROUP", splitPoints); final List<Integer> partitionIds1 = partitioner.getPartitionIds(new Object[]{0L}); final List<Integer> partitionIds2 = partitioner.getPartitionIds(new Object[]{1L}); final List<Integer> partitionIds3 = partitioner.getPartitionIds(new Object[]{2L}); final List<Integer> partitionIds4 = partitioner.getPartitionIds(new Object[]{6L}); final List<Integer> partitionIds5 = partitioner.getPartitionIds(new Object[]{1L, 2L}); final List<Integer> partitionIds6 = partitioner.getPartitionIds(new Object[]{1L, 4L}); final List<Integer> partitionIds7 = partitioner.getPartitionIds(new Object[]{1L, 6L}); final List<Integer> partitionIds8 = partitioner.getPartitionIds(new Object[]{1L, 7L}); final List<Integer> partitionIds9 = partitioner.getPartitionIds(new Object[]{1L, 10L}); final List<Integer> partitionIds10 = partitioner.getPartitionIds(new Object[]{2L, 10L}); final List<Integer> partitionIds11 = partitioner.getPartitionIds(new Object[]{5L, 3L}); final List<Integer> partitionIds12 = partitioner.getPartitionIds(new Object[]{5L, 4L}); final List<Integer> partitionIds13 = partitioner.getPartitionIds(new Object[]{10L, 10L}); assertEquals(Arrays.asList(0), partitionIds1); assertEquals(Arrays.asList(0, 1, 2, 3), partitionIds2); assertEquals(Arrays.asList(3), partitionIds3); assertEquals(Arrays.asList(4), partitionIds4); assertEquals(Arrays.asList(0), partitionIds5); assertEquals(Arrays.asList(1), partitionIds6); assertEquals(Arrays.asList(1), partitionIds7); assertEquals(Arrays.asList(2), partitionIds8); assertEquals(Arrays.asList(3), partitionIds9); assertEquals(Arrays.asList(3), partitionIds10); assertEquals(Arrays.asList(3), partitionIds11); assertEquals(Arrays.asList(4), partitionIds12); assertEquals(Arrays.asList(4), partitionIds13); }
@Test public void testGetPartitionIdsEmptyGroupPartitioner() { final GroupPartitioner partitioner = new GroupPartitioner("GROUP", new ArrayList<>()); final List<Integer> partitionIds = partitioner.getPartitionIds(new Object[]{1L}); assertEquals(Arrays.asList(0), partitionIds); }