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

  private var rangeBounds: Array[K] = {

    if (partitions <= 1) {

      Array.empty

    } else {

      // This is the sample size we need to have roughly balanced output partitions, capped at 1M.

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

      // Assume the input partitions are roughly balanced and over-sample a little bit.

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

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

      if (numItems == 0L) {

        Array.empty

      } else {

        // If a partition contains much more than the average number of items, we re-sample from it

        // to ensure that enough items are collected from that partition.

        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.length).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))

        }

        RangePartitioner.determineBounds(candidates, partitions)

      }

    }

  }

private[spark] object RangePartitioner {

  /**

   * Sketches the input RDD via reservoir sampling on each partition.

   *

   * @param rdd the input RDD to sketch

   * @param sampleSizePerPartition max sample size per partition

   * @return (total number of items, an array of (partitionId, number of items, sample))

   */

  def sketch[K : ClassTag](

      rdd: RDD[K],

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

    val shift = rdd.id

    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).sum

    (numItems, sketched)

  }

/**

   * Determines the bounds for range partitioning from candidates with weights indicating how many

   * items each represents. Usually this is 1 over the probability used to sample this candidate.

   *

   * @param candidates unordered candidates with weights

   * @param partitions number of partitions

   * @return selected bounds

   */

  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 sortByKey(ascending: Boolean = true, numPartitions: Int = self.partitions.length): RDD[(K, V)] = self.withScope

  {

    val part = new RangePartitioner(numPartitions, self, ascending)

    new ShuffledRDD[K, V, V](self, part)

      .setKeyOrdering(if (ascending) ordering else ordering.reverse)

  }