diff options
author | Matei Zaharia <matei@eecs.berkeley.edu> | 2012-03-17 13:08:36 -0700 |
---|---|---|
committer | Matei Zaharia <matei@eecs.berkeley.edu> | 2012-03-17 13:08:36 -0700 |
commit | c7af538ac160727147eea6a1008dc16a2efd802e (patch) | |
tree | 6bad957c7ac56b21483406161304e59d22cd7d82 /core | |
parent | a5e2b6a6bd682eac2859b9105feabaf95ec1d059 (diff) | |
download | spark-c7af538ac160727147eea6a1008dc16a2efd802e.tar.gz spark-c7af538ac160727147eea6a1008dc16a2efd802e.tar.bz2 spark-c7af538ac160727147eea6a1008dc16a2efd802e.zip |
Some fixes to sorting for when the RDD has fewer elements than the
number of partitions we ask to partition it into. Also, removed a test
that was taking way too long to run.
Diffstat (limited to 'core')
-rw-r--r-- | core/src/main/scala/spark/Partitioner.scala | 55 | ||||
-rw-r--r-- | core/src/test/scala/spark/SortingSuite.scala | 15 |
2 files changed, 47 insertions, 23 deletions
diff --git a/core/src/main/scala/spark/Partitioner.scala b/core/src/main/scala/spark/Partitioner.scala index d1e65edcda..ac61fe3b54 100644 --- a/core/src/main/scala/spark/Partitioner.scala +++ b/core/src/main/scala/spark/Partitioner.scala @@ -25,34 +25,49 @@ class HashPartitioner(partitions: Int) extends Partitioner { } } -class RangePartitioner[K <% Ordered[K],V](partitions: Int, rdd: RDD[(K,V)], ascending: Boolean = true) +class RangePartitioner[K <% Ordered[K]: ClassManifest, V]( + partitions: Int, rdd: RDD[(K,V)], + ascending: Boolean = true) extends Partitioner { - def numPartitions = partitions + private val rangeBounds: Array[K] = { + val rddSize = rdd.count() + val maxSampleSize = partitions * 10.0 + val frac = math.min(maxSampleSize / math.max(rddSize, 1), 1.0) + val rddSample = rdd.sample(true, frac, 1).map(_._1).collect() + .sortWith((x, y) => if (ascending) x < y else x > y) + if (rddSample.length == 0) { + Array() + } else { + val bounds = new Array[K](partitions) + for (i <- 0 until partitions) { + bounds(i) = rddSample(i * rddSample.length / partitions) + } + bounds + } + } + + def numPartitions = rangeBounds.length - val rddSize = rdd.count() - val maxSampleSize = partitions*10.0 - val frac = 1.0.min(maxSampleSize / rddSize) - val rddSample = rdd.sample(true, frac, 1).collect.toList - .sortWith((x, y) => if (ascending) x._1 < y._1 else x._1 > y._1) - .map(_._1) - val bucketSize = rddSample.size / partitions - val rangeBounds = rddSample.zipWithIndex.filter(_._2 % bucketSize == 0) - .map(_._1).slice(1, partitions) - - def getPartition(key: Any): Int = { + def getPartition(key: Any): Int = { + // TODO: Use a binary search here if number of partitions is large val k = key.asInstanceOf[K] - val p = rangeBounds.zipWithIndex.foldLeft(0) { - case (part, (bound, index)) => - if (k > bound) index + 1 else part - } - if (ascending) p else numPartitions-1-p + var partition = 0 + while (partition < rangeBounds.length - 1 && k > rangeBounds(partition)) { + partition += 1 + } + if (ascending) { + partition + } else { + rangeBounds.length - 1 - partition + } } override def equals(other: Any): Boolean = other match { case r: RangePartitioner[_,_] => - r.numPartitions == numPartitions & r.rangeBounds == rangeBounds - case _ => false + r.rangeBounds.sameElements(rangeBounds) + case _ => + false } } diff --git a/core/src/test/scala/spark/SortingSuite.scala b/core/src/test/scala/spark/SortingSuite.scala index 0d1f229eda..caff884966 100644 --- a/core/src/test/scala/spark/SortingSuite.scala +++ b/core/src/test/scala/spark/SortingSuite.scala @@ -29,11 +29,20 @@ class SortingSuite extends FunSuite { sc.stop() } - test("sortHighParallelism") { + test("morePartitionsThanElements") { val sc = new SparkContext("local", "test") val rand = new scala.util.Random() - val pairArr = Array.fill(3000) { (rand.nextInt(), rand.nextInt()) } - val pairs = sc.parallelize(pairArr, 300) + val pairArr = Array.fill(10) { (rand.nextInt(), rand.nextInt()) } + val pairs = sc.parallelize(pairArr, 30) + assert(pairs.sortByKey().collect() === pairArr.sortBy(_._1)) + sc.stop() + } + + test("emptyRDD") { + val sc = new SparkContext("local", "test") + val rand = new scala.util.Random() + val pairArr = new Array[(Int, Int)](0) + val pairs = sc.parallelize(pairArr) assert(pairs.sortByKey().collect() === pairArr.sortBy(_._1)) sc.stop() } |