package com.tzld.piaoquan.ad.spark.common
import scala.util.Random
import org.apache.spark.rdd.{PartitionCoalescer, PartitionGroup, RDD}

class RandomUniformCoalescer extends PartitionCoalescer with Serializable  {
  override def coalesce(maxPartitions: Int, parent: RDD[_]): Array[PartitionGroup] = {
    val parentPartitions = parent.getNumPartitions
    val shufflePartitions = Random.shuffle(parent.partitions.toList).toArray
    val expectNumPartitions = math.min(maxPartitions, parentPartitions)
    // 避免出现最后一个group过大导致长尾现象，故groupSize放缩至更大的值
    val groupSize = (parentPartitions + expectNumPartitions - 1) / expectNumPartitions
    val numPartitions = (parentPartitions + groupSize - 1) / groupSize
    val groups = new Array[PartitionGroup](numPartitions)
    for (i <- 0 until numPartitions) {
      val start = i * groupSize
      val end = math.min(start + groupSize, parentPartitions)
      groups(i) = new PartitionGroup()
      groups(i).partitions.appendAll(shufflePartitions.slice(start, end))
    }
    groups
  }
}
