AbstractKuduInputPartitioner implements Partitioner<AbstractKuduInputOperator> { @Override public Collection<Partition<AbstractKuduInputOperator>> definePartitions( Collection<Partition<AbstractKuduInputOperator>> partitions, PartitioningContext context) { if ( partitions != null) { LOG.info("The current partitioner plan has " + partitions.size() + " operators before redefining"); } List<Partition<AbstractKuduInputOperator>> partitionsForInputOperator = new ArrayList<>(); List<KuduPartitionScanAssignmentMeta> assignmentMetaList = new ArrayList<>(); try { assignmentMetaList.addAll(getListOfPartitionAssignments(partitions,context)); } catch (Exception e) { throw new RuntimeException("Aborting partition planning as Kudu meta data could not be obtained", e); } LOG.info("Maximum possible Kudu input operator partition count is " + assignmentMetaList.size()); Map<Integer,List<KuduPartitionScanAssignmentMeta>> assignments = assign(assignmentMetaList,context); boolean requiresRepartitioning = false; if ( partitions == null) { requiresRepartitioning = true; } else { for ( Partition<AbstractKuduInputOperator> aPartition : partitions) { if ( !aPartition.getPartitionedInstance().isPartitioned()) { requiresRepartitioning = true; break; } } } if ( requiresRepartitioning ) { partitions.clear(); LOG.info("Clearing all of the current partitions and setting up new ones"); partitions.clear(); for ( int i = 0; i < assignments.size(); i++) { List<KuduPartitionScanAssignmentMeta> assignmentForThisOperator = assignments.get(i); partitionsForInputOperator.add(clonePartitionAndAssignScanMeta(assignmentForThisOperator)); LOG.info("Assigned apex operator " + i + " with " + assignmentForThisOperator.size() + " kudu mappings"); } LOG.info("Returning " + partitionsForInputOperator.size() + " partitions for the input operator"); return partitionsForInputOperator; } else { LOG.info("Not making any changes to the partitions"); return partitions; } } AbstractKuduInputPartitioner(AbstractKuduInputOperator prototypeOperator); List<KuduScanToken> getKuduScanTokensForSelectAllColumns(); int getNumberOfPartitions(PartitioningContext context); List<KuduPartitionScanAssignmentMeta> getListOfPartitionAssignments(
      Collection<Partition<AbstractKuduInputOperator>> partitions, PartitioningContext context); Partitioner.Partition<AbstractKuduInputOperator> clonePartitionAndAssignScanMeta(
      List<KuduPartitionScanAssignmentMeta> scanAssignmentsForThisPartition); @Override Collection<Partition<AbstractKuduInputOperator>> definePartitions(
      Collection<Partition<AbstractKuduInputOperator>> partitions, PartitioningContext context); @Override void partitioned(Map<Integer, Partition<AbstractKuduInputOperator>> partitions); abstract Map<Integer, List<KuduPartitionScanAssignmentMeta>> assign(
      List<KuduPartitionScanAssignmentMeta> totalList, PartitioningContext context); AbstractKuduInputOperator getPrototypeKuduInputOperator(); void setPrototypeKuduInputOperator(AbstractKuduInputOperator prototypeKuduInputOperator);  }
@KuduClusterTestContext(kuduClusterBasedTest = true) @Test public void testDefinePartitions() throws Exception { numberOfKuduInputOperatorPartitions = 7; initOperatorState(); AbstractKuduInputPartitioner partitioner = unitTestStepwiseScanInputOperator.getPartitioner(); assertEquals(KuduClientTestCommons.TOTAL_KUDU_TABLETS_FOR_UNITTEST_TABLE,partitioner.definePartitions( new ArrayList<Partitioner.Partition<AbstractKuduInputOperator>>(),partitioningContext).size()); partitonScanStrategy = KuduPartitionScanStrategy.MANY_TABLETS_PER_OPERATOR; initOperatorState(); partitioner = unitTestStepwiseScanInputOperator.getPartitioner(); Collection<Partitioner.Partition<AbstractKuduInputOperator>> partitionsCalculated = partitioner.definePartitions( new ArrayList<Partitioner.Partition<AbstractKuduInputOperator>>(),partitioningContext); assertEquals(7,partitionsCalculated.size()); int maxPartitionPerOperator = -1; int minPartitionPerOperator = Integer.MAX_VALUE; Iterator<Partitioner.Partition<AbstractKuduInputOperator>> iteratorForPartitionScan = partitionsCalculated.iterator(); while ( iteratorForPartitionScan.hasNext()) { Partitioner.Partition<AbstractKuduInputOperator> anOperatorPartition = iteratorForPartitionScan.next(); AbstractKuduInputOperator<UnitTestTablePojo,InputOperatorControlTuple> anOperatorHandle = anOperatorPartition.getPartitionedInstance(); List<KuduPartitionScanAssignmentMeta> partitionsAssigned = anOperatorHandle.getPartitionPieAssignment(); if ( partitionsAssigned.size() > maxPartitionPerOperator ) { maxPartitionPerOperator = partitionsAssigned.size(); } if ( partitionsAssigned.size() < minPartitionPerOperator) { minPartitionPerOperator = partitionsAssigned.size(); } } assertEquals(2,maxPartitionPerOperator); assertEquals(1,minPartitionPerOperator); numberOfKuduInputOperatorPartitions = 5; partitonScanStrategy = KuduPartitionScanStrategy.ONE_TABLET_PER_OPERATOR; }