docs: public static <T> JavaRDD<T> repartitionBalanceIfRequired(JavaRDD<T> rdd, Repartition repartition,
                      int objectsPerPartition, int numPartitions) {
          int origNumPartitions = rdd.partitions().size();
          switch (repartition) {
              case Never:
                  return rdd;
              case NumPartitionsWorkersDiffers:
                  if (origNumPartitions == numPartitions)
                      return rdd;
              case Always:
                  //Repartition: either always, or origNumPartitions != numWorkers
  
                  //First: count number of elements in each partition. Need to know this so we can work out how to properly index each example,
                  // so we can in turn create properly balanced partitions after repartitioning
                  //Because the objects (DataSets etc) should be small, this should be OK
  
                  //Count each partition...
                  List<Tuple2<Integer, Integer>> partitionCounts =
                                  rdd.mapPartitionsWithIndex(new CountPartitionsFunction<T>(), true).collect();
                  int totalObjects = 0;
                  int initialPartitions = partitionCounts.size();
  
                  boolean allCorrectSize = true;
                  int x = 0;
                  for (Tuple2<Integer, Integer> t2 : partitionCounts) {
                      int partitionSize = t2._2();
                      allCorrectSize &= (partitionSize == objectsPerPartition);
                      totalObjects += t2._2();
                  }
  
                  if (numPartitions * objectsPerPartition < totalObjects) {
                      allCorrectSize = true;
                      for (Tuple2<Integer, Integer> t2 : partitionCounts) {
                          allCorrectSize &= (t2._2() == objectsPerPartition);
                      }
                  }
  
                  if (initialPartitions == numPartitions && allCorrectSize) {
                      //Don't need to do any repartitioning here - already in the format we want
                      return rdd;
                  }
  
                  //Index each element for repartitioning (can only do manual repartitioning on a JavaPairRDD)
                  JavaPairRDD<Integer, T> pairIndexed = indexedRDD(rdd);
  
                  int remainder = (totalObjects - numPartitions * objectsPerPartition) % numPartitions;
                  log.info("Amount to rebalance: numPartitions={}, objectsPerPartition={}, remainder={}", numPartitions, objectsPerPartition, remainder);
                  pairIndexed = pairIndexed
                                  .partitionBy(new BalancedPartitioner(numPartitions, objectsPerPartition, remainder));
                  return pairIndexed.values();
              default:
                  throw new RuntimeException("Unknown setting for repartition: " + repartition);
          }
      }