[GitHub] [spark] mridulm commented on a change in pull request #31480: [SPARK-32384][CORE] repartitionAndSortWithinPartitions avoid shuffle with same partitioner

[GitBox <gi...@apache.org>]

mridulm commented on a change in pull request #31480:
URL: https://github.com/apache/spark/pull/31480#discussion_r593416332



##########
File path: core/src/main/scala/org/apache/spark/rdd/OrderedRDDFunctions.scala
##########
@@ -73,7 +75,23 @@ class OrderedRDDFunctions[K : Ordering : ClassTag,
    * because it can push the sorting down into the shuffle machinery.
    */
   def repartitionAndSortWithinPartitions(partitioner: Partitioner): RDD[(K, V)] = self.withScope {
-    new ShuffledRDD[K, V, V](self, partitioner).setKeyOrdering(ordering)
+    if (self.partitioner == Some(partitioner)) {
+      self.mapPartitions(iter => {
+        val context = TaskContext.get
+        val sorter = new ExternalSorter[K, V, V](context, None, None, Some(ordering))
+        sorter.insertAll(iter)
+        context.taskMetrics.incMemoryBytesSpilled(sorter.memoryBytesSpilled)
+        context.taskMetrics.incDiskBytesSpilled(sorter.diskBytesSpilled)
+        context.taskMetrics.incPeakExecutionMemory(sorter.peakMemoryUsedBytes)
+        // Use completion callback to stop sorter if task was finished/cancelled.
+        context.addTaskCompletionListener[Unit](_ => sorter.stop)
+        val outputIter = new InterruptibleIterator(context,
+          sorter.iterator.asInstanceOf[Iterator[(K, V)]])
+        CompletionIterator[(K, V), Iterator[(K, V)]](outputIter, sorter.stop)
+      }, preservesPartitioning = true)
+    } else {

Review comment:
       Nice fix !

##########
File path: core/src/test/scala/org/apache/spark/rdd/RDDSuite.scala
##########
@@ -860,6 +860,32 @@ class RDDSuite extends SparkFunSuite with SharedSparkContext with Eventually {
     assert(partitions(1) === Seq((1, 3), (3, 8), (3, 8)))
   }
 
+  test("SPARK-32384: repartitionAndSortWithinPartitions without shuffle") {
+    val data = sc.parallelize(Seq((0, 5), (3, 8), (2, 6), (0, 8), (3, 8), (1, 3)), 2)
+
+    class ModPartitioner(val numPartitions: Int) extends Partitioner {
+      def getPartition(key: Any): Int = key.asInstanceOf[Int] % numPartitions
+
+      override def equals(other: Any): Boolean = other match {
+        case h: ModPartitioner => h.numPartitions == this.numPartitions
+        case _ => false
+      }
+
+      override def hashCode: Int = numPartitions
+    }

Review comment:
       Remove this and use `HashPartitioner` ?

##########
File path: core/src/test/scala/org/apache/spark/rdd/RDDSuite.scala
##########
@@ -860,6 +860,32 @@ class RDDSuite extends SparkFunSuite with SharedSparkContext with Eventually {
     assert(partitions(1) === Seq((1, 3), (3, 8), (3, 8)))
   }
 
+  test("SPARK-32384: repartitionAndSortWithinPartitions without shuffle") {
+    val data = sc.parallelize(Seq((0, 5), (3, 8), (2, 6), (0, 8), (3, 8), (1, 3)), 2)
+
+    class ModPartitioner(val numPartitions: Int) extends Partitioner {
+      def getPartition(key: Any): Int = key.asInstanceOf[Int] % numPartitions
+
+      override def equals(other: Any): Boolean = other match {
+        case h: ModPartitioner => h.numPartitions == this.numPartitions
+        case _ => false
+      }
+
+      override def hashCode: Int = numPartitions
+    }
+
+    val partitioner = new ModPartitioner(2)
+    val agged = data.reduceByKey(partitioner, _ + _)
+    assert(agged.partitioner == Some(partitioner))
+
+    val sorted = agged.repartitionAndSortWithinPartitions(partitioner)
+    assert(sorted.partitioner == Some(partitioner))
+
+    val partitions = sorted.glom().collect()
+    assert(partitions(0) === Seq((0, 13), (2, 6)))
+    assert(partitions(1) === Seq((1, 3), (3, 16)))

Review comment:
       This test is not testing if a shuffle was avoided - just that `repartitionAndSortWithinPartitions` worked, which is already tested elsewhere.
   You will have to test to see how many stages are executed - to validate `repartitionAndSortWithinPartitions` became a narrow dependency instead of shuffle dependency




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org