KuduOneToManyPartitioner extends AbstractKuduInputPartitioner { @Override public Map<Integer, List<KuduPartitionScanAssignmentMeta>> assign(List<KuduPartitionScanAssignmentMeta> totalList, PartitioningContext context) { Map<Integer,List<KuduPartitionScanAssignmentMeta>> partitionAssignments = new HashMap<>(); int partitionCount = getNumberOfPartitions(context); if ( partitionCount <= 0) { LOG.error(" Partition count cannot be zero "); partitionCount = 1; } int idealDistributionRatio = (totalList.size() / partitionCount) + 1; LOG.info(" Distributing not more than " + idealDistributionRatio + " partitions per input operator"); int counterForLoopingTotal = 0; int totalSizeOfKuduScanAssignments = totalList.size(); for ( int i = 0; i < partitionCount; i++) { partitionAssignments.put(i, new ArrayList<KuduPartitionScanAssignmentMeta>()); } while ( counterForLoopingTotal < totalSizeOfKuduScanAssignments) { for (int i = 0; i < partitionCount; i++) { List<KuduPartitionScanAssignmentMeta> assignmentsForThisOperatorId = partitionAssignments.get(i); if (counterForLoopingTotal < totalSizeOfKuduScanAssignments) { assignmentsForThisOperatorId.add(totalList.get(counterForLoopingTotal)); counterForLoopingTotal += 1; } else { break; } } } return partitionAssignments; } KuduOneToManyPartitioner(AbstractKuduInputOperator prototypeOperator); @Override Map<Integer, List<KuduPartitionScanAssignmentMeta>> assign(List<KuduPartitionScanAssignmentMeta> totalList,
      PartitioningContext context);  }
@KuduClusterTestContext(kuduClusterBasedTest = false) @Test public void testAssignPartitions() throws Exception { AbstractKuduInputOperator<UnitTestTablePojo,InputOperatorControlTuple> mockedInputOperator = PowerMockito.mock(AbstractKuduInputOperator.class); when(mockedInputOperator.getNumberOfPartitions()).thenReturn(5); PowerMockito.mockStatic(KryoCloneUtils.class); when(KryoCloneUtils.cloneObject(mockedInputOperator)).thenReturn(mockedInputOperator); KuduOneToManyPartitioner kuduOneToManyPartitioner = new KuduOneToManyPartitioner(mockedInputOperator); buildMockWiring(mockedInputOperator, KuduClientTestCommons.TOTAL_KUDU_TABLETS_FOR_UNITTEST_TABLE); kuduOneToManyPartitioner.setPrototypeKuduInputOperator(mockedInputOperator); Map<Integer,List<KuduPartitionScanAssignmentMeta>> assignedPartitions = kuduOneToManyPartitioner.assign( kuduOneToManyPartitioner.getListOfPartitionAssignments( new ArrayList<Partitioner.Partition<AbstractKuduInputOperator>>(), partitioningContext),partitioningContext); assertThat(assignedPartitions.size(), is(5)); for (List<KuduPartitionScanAssignmentMeta> eachOperatorassignment: assignedPartitions.values()) { assertThat(eachOperatorassignment.size(), lessThanOrEqualTo(3)); } for (List<KuduPartitionScanAssignmentMeta> eachOperatorassignment: assignedPartitions.values()) { assertThat(eachOperatorassignment.size(), greaterThanOrEqualTo(2)); } }