class HashPartitioner(partitions: Int) extends Partitioner {

??require(partitions >= 0, s"Number of partitions ($partitions) cannot be negative.")

?

??def numPartitions: Int = partitions

?

??def getPartition(key: Any): Int = key match {

????case null => 0

????case _ => Utils.nonNegativeMod(key.hashCode, numPartitions)

??}

?

??override def equals(other: Any): Boolean = other match {

????case h: HashPartitioner =>

??????h.numPartitions == numPartitions

????case _ =>

??????false

??}

?

??override def hashCode: Int = numPartitions

}

// An array of upper bounds for the first (partitions - 1) partitions

private val rangeBounds: Array[K] = {

????if (partitions == 1) {

??????Array()

????} else {

??????val rddSize = rdd.count()

??????val maxSampleSize = partitions * 20.0

??????val frac = math.min(maxSampleSize / math.max(rddSize, 1), 1.0)

??????val rddSample = rdd.sample(false, frac, 1).map(_._1).collect().sorted

??????if (rddSample.length == 0) {

????????Array()

??????} else {

????????val bounds = new Array[K](partitions - 1)

????????for (i <- 0 until partitions - 1) {

??????????val index = (rddSample.length - 1) * (i + 1) / partitions

??????????bounds(i) = rddSample(index)

????????}

????????bounds

??????}

????}

}

val sampleSize = math.min(20.0 * partitions, 1e6)

val sampleSizePerPartition = math.ceil(3.0 * sampleSize / rdd.partitions.size).toInt

?

val (numItems, sketched) = RangePartitioner.sketch(rdd.map(_._1), sampleSizePerPartition)

?

def sketch[K : ClassTag](

??????rdd: RDD[K],

??????sampleSizePerPartition: Int): (Long, Array[(Int, Int, Array[K])]) = {

????val shift = rdd.id

????// val classTagK = classTag[K] // to avoid serializing the entire partitioner object

????val sketched = rdd.mapPartitionsWithIndex { (idx, iter) =>

??????val seed = byteswap32(idx ^ (shift << 16))

??????val (sample, n) = SamplingUtils.reservoirSampleAndCount(

????????iter, sampleSizePerPartition, seed)

??????Iterator((idx, n, sample))

????}.collect()

????val numItems = sketched.map(_._2.toLong).sum

????(numItems, sketched)

}

?

def reservoirSampleAndCount[T: ClassTag](

??????input: Iterator[T],

??????k: Int,

??????seed: Long = Random.nextLong())

????: (Array[T], Int) = {

????val reservoir = new Array[T](k)

????// Put the first k elements in the reservoir.

????var i = 0

????while (i < k && input.hasNext) {

??????val item = input.next()

??????reservoir(i) = item

??????i += 1

????}

?

????// If we have consumed all the elements, return them. Otherwise do the replacement.

????if (i < k) {

??????// If input size < k, trim the array to return only an array of input size.

??????val trimReservoir = new Array[T](i)

??????System.arraycopy(reservoir, 0, trimReservoir, 0, i)

??????(trimReservoir, i)

????} else {

??????// If input size > k, continue the sampling process.

??????val rand = new XORShiftRandom(seed)

??????while (input.hasNext) {

????????val item = input.next()

????????val replacementIndex = rand.nextInt(i)

????????if (replacementIndex < k) {

??????????reservoir(replacementIndex) = item

????????}

????????i += 1

??????}

??????(reservoir, i)

}

}

val fraction = math.min(sampleSize / math.max(numItems, 1L), 1.0)

val candidates = ArrayBuffer.empty[(K, Float)]

val imbalancedPartitions = mutable.Set.empty[Int]

sketched.foreach { case (idx, n, sample) =>

??if (fraction * n > sampleSizePerPartition) {

????imbalancedPartitions += idx

??} else {

????// The weight is 1 over the sampling probability.

????val weight = (n.toDouble / sample.size).toFloat

????for (key <- sample) {

??????candidates += ((key, weight))

????}

??}

}

if (imbalancedPartitions.nonEmpty) {

??// Re-sample imbalanced partitions with the desired sampling probability.

??val imbalanced = new PartitionPruningRDD(rdd.map(_._1), imbalancedPartitions.contains)

??val seed = byteswap32(-rdd.id - 1)

??val reSampled = imbalanced.sample(withReplacement = false, fraction, seed).collect()

??val weight = (1.0 / fraction).toFloat

??candidates ++= reSampled.map(x => (x, weight))

}

def determineBounds[K : Ordering : ClassTag](

????candidates: ArrayBuffer[(K, Float)],

????partitions: Int): Array[K] = {

??val ordering = implicitly[Ordering[K]]

??val ordered = candidates.sortBy(_._1)

??val numCandidates = ordered.size

??val sumWeights = ordered.map(_._2.toDouble).sum

??val step = sumWeights / partitions

??var cumWeight = 0.0

??var target = step

??val bounds = ArrayBuffer.empty[K]

??var i = 0

??var j = 0

??var previousBound = Option.empty[K]

??while ((i < numCandidates) && (j < partitions - 1)) {

????val (key, weight) = ordered(i)

????cumWeight += weight

????if (cumWeight > target) {

??????// Skip duplicate values.

??????if (previousBound.isEmpty || ordering.gt(key, previousBound.get)) {

????????bounds += key

????????target += step

????????j += 1

????????previousBound = Some(key)

??????}

????}

????i += 1

??}

??bounds.toArray

}

def getPartition(key: Any): Int = {

??val k = key.asInstanceOf[K]

??var partition = 0

??if (rangeBounds.length <= 128) {

????// If we have less than 128 partitions naive search

????while (partition < rangeBounds.length && ordering.gt(k, rangeBounds(partition))) {

??????partition += 1

????}

??} else {

????// Determine which binary search method to use only once.

????partition = binarySearch(rangeBounds, k)

????// binarySearch either returns the match location or -[insertion point]-1

????if (partition < 0) {

??????partition = -partition-1

????}

????if (partition > rangeBounds.length) {

??????partition = rangeBounds.length

????}

??}

??if (ascending) {

????partition

??} else {

????rangeBounds.length - partition

??}

}