Move sort into shuffle implementations

Author: jerryshao

core/src/main/scala/org/apache/spark/Dependency.scala

@@ -62,9 +62,12 @@ class ShuffleDependency[K, V, C](
     val serializer: Option[Serializer] = None,
     val keyOrdering: Option[Ordering[K]] = None,
     val aggregator: Option[Aggregator[K, V, C]] = None,
-    val mapSideCombine: Boolean = false)
+    val mapSideCombine: Boolean = false,
+    val ascending: Boolean = true)
   extends Dependency(rdd.asInstanceOf[RDD[Product2[K, V]]]) {
 
+  def isKeySorted = keyOrdering.isDefined
+
   val shuffleId: Int = rdd.context.newShuffleId()
 
   val shuffleHandle: ShuffleHandle = rdd.context.env.shuffleManager.registerShuffle(

core/src/main/scala/org/apache/spark/rdd/OrderedRDDFunctions.scala

@@ -57,14 +57,8 @@ class OrderedRDDFunctions[K : Ordering : ClassTag,
    */
   def sortByKey(ascending: Boolean = true, numPartitions: Int = self.partitions.size): RDD[P] = {
     val part = new RangePartitioner(numPartitions, self, ascending)
-    val shuffled = new ShuffledRDD[K, V, V, P](self, part).setKeyOrdering(ordering)
-    shuffled.mapPartitions(iter => {
-      val buf = iter.toArray
-      if (ascending) {
-        buf.sortWith((x, y) => ordering.lt(x._1, y._1)).iterator
-      } else {
-        buf.sortWith((x, y) => ordering.gt(x._1, y._1)).iterator
-      }
-    }, preservesPartitioning = true)
+    new ShuffledRDD[K, V, V, P](self, part)
+      .setKeyOrdering(ordering)
+      .setAscendingFlag(ascending)
   }
 }

core/src/main/scala/org/apache/spark/rdd/ShuffledRDD.scala

@@ -51,6 +51,8 @@ class ShuffledRDD[K, V, C, P <: Product2[K, C] : ClassTag](
 
   private var mapSideCombine: Boolean = false
 
+  private var ascending: Boolean = true
+
   /** Set a serializer for this RDD's shuffle, or null to use the default (spark.serializer) */
   def setSerializer(serializer: Serializer): ShuffledRDD[K, V, C, P] = {
     this.serializer = Option(serializer)
@@ -63,6 +65,11 @@ class ShuffledRDD[K, V, C, P <: Product2[K, C] : ClassTag](
     this
   }
 
+  def setAscendingFlag(ascending: Boolean): ShuffledRDD[K, V, C, P] = {
+    this.ascending = ascending
+    this
+  }
+
   /** Set aggregator for RDD's shuffle. */
   def setAggregator(aggregator: Aggregator[K, V, C]): ShuffledRDD[K, V, C, P] = {
     this.aggregator = Option(aggregator)
@@ -76,7 +83,8 @@ class ShuffledRDD[K, V, C, P <: Product2[K, C] : ClassTag](
   }
 
   override def getDependencies: Seq[Dependency[_]] = {
-    List(new ShuffleDependency(prev, part, serializer, keyOrdering, aggregator, mapSideCombine))
+    List(new ShuffleDependency(prev, part, serializer,
+      keyOrdering, aggregator, mapSideCombine, ascending))
   }
 
   override val partitioner = Some(part)

core/src/main/scala/org/apache/spark/shuffle/hash/HashShuffleReader.scala

@@ -38,7 +38,7 @@ class HashShuffleReader[K, C](
     val iter = BlockStoreShuffleFetcher.fetch(handle.shuffleId, startPartition, context,
       Serializer.getSerializer(dep.serializer))
 
-    if (dep.aggregator.isDefined) {
+    val aggregatedIter = if (dep.aggregator.isDefined) {
       if (dep.mapSideCombine) {
         new InterruptibleIterator(context, dep.aggregator.get.combineCombinersByKey(iter, context))
       } else {
@@ -49,6 +49,17 @@ class HashShuffleReader[K, C](
     } else {
       iter
     }
+
+    dep.keyOrdering.map { ordering =>
+      val buf = aggregatedIter.toArray
+      if (dep.ascending) {
+        buf.sortWith((x, y) => ordering.lt(x._1, y._1)).iterator
+      } else {
+        buf.sortWith((x, y) => ordering.gt(x._1, y._1)).iterator
+      }
+    }.getOrElse {
+      aggregatedIter
+    }
   }
 
   /** Close this reader */