overload getReader

Author: Ngone51

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

@@ -145,8 +145,7 @@ class CoGroupedRDD[K: ClassTag](
         val metrics = context.taskMetrics().createTempShuffleReadMetrics()
         val it = SparkEnv.get.shuffleManager
           .getReader(
-            shuffleDependency.shuffleHandle, 0, Int.MaxValue,
-            split.index, split.index + 1, context, metrics)
+            shuffleDependency.shuffleHandle, split.index, split.index + 1, context, metrics)
           .read()
         rddIterators += ((it, depNum))
     }

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

@@ -103,8 +103,7 @@ class ShuffledRDD[K: ClassTag, V: ClassTag, C: ClassTag](
     val dep = dependencies.head.asInstanceOf[ShuffleDependency[K, V, C]]
     val metrics = context.taskMetrics().createTempShuffleReadMetrics()
     SparkEnv.get.shuffleManager.getReader(
-      dep.shuffleHandle, 0, Int.MaxValue,
-      split.index, split.index + 1, context, metrics)
+      dep.shuffleHandle, split.index, split.index + 1, context, metrics)
       .read()
       .asInstanceOf[Iterator[(K, C)]]
   }

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

@@ -105,8 +105,6 @@ private[spark] class SubtractedRDD[K: ClassTag, V: ClassTag, W: ClassTag](
           val iter = SparkEnv.get.shuffleManager
             .getReader(
               shuffleDependency.shuffleHandle,
-              0,
-              Int.MaxValue,
               partition.index,
               partition.index + 1,
               context,

core/src/main/scala/org/apache/spark/shuffle/ShuffleManager.scala

@@ -43,6 +43,22 @@ private[spark] trait ShuffleManager {
       context: TaskContext,
       metrics: ShuffleWriteMetricsReporter): ShuffleWriter[K, V]
 
+
+  /**
+   * Get a reader for a range of reduce partitions (startPartition to endPartition-1, inclusive) to
+   * read from all map outputs of the shuffle.
+   *
+   * Called on executors by reduce tasks.
+   */
+  final def getReader[K, C](
+      handle: ShuffleHandle,
+      startPartition: Int,
+      endPartition: Int,
+      context: TaskContext,
+      metrics: ShuffleReadMetricsReporter): ShuffleReader[K, C] = {
+    getReader(handle, 0, Int.MaxValue, startPartition, endPartition, context, metrics)
+  }
+
   /**
    * Get a reader for a range of reduce partitions (startPartition to endPartition-1, inclusive) to
    * read from a range of map outputs(startMapIndex to endMapIndex-1, inclusive).

core/src/test/scala/org/apache/spark/ShuffleSuite.scala

@@ -411,8 +411,7 @@ abstract class ShuffleSuite extends SparkFunSuite with Matchers with LocalSparkC
     val taskContext = new TaskContextImpl(
       1, 0, 0, 2L, 0, taskMemoryManager, new Properties, metricsSystem)
     val metrics = taskContext.taskMetrics.createTempShuffleReadMetrics()
-    val reader = manager.getReader[Int, Int](shuffleHandle, 0, Int.MaxValue,
-      0, 1, taskContext, metrics)
+    val reader = manager.getReader[Int, Int](shuffleHandle, 0, 1, taskContext, metrics)
     TaskContext.unset()
     val readData = reader.read().toIndexedSeq
     assert(readData === data1.toIndexedSeq || readData === data2.toIndexedSeq)

core/src/test/scala/org/apache/spark/scheduler/CustomShuffledRDD.scala

@@ -106,8 +106,7 @@ class CustomShuffledRDD[K, V, C](
     val part = p.asInstanceOf[CustomShuffledRDDPartition]
     val metrics = context.taskMetrics().createTempShuffleReadMetrics()
     SparkEnv.get.shuffleManager.getReader(
-      dependency.shuffleHandle, 0, Int.MaxValue, part.startIndexInParent,
-      part.endIndexInParent, context, metrics)
+      dependency.shuffleHandle, part.startIndexInParent, part.endIndexInParent, context, metrics)
       .read()
       .asInstanceOf[Iterator[(K, C)]]
   }

sql/core/src/main/scala/org/apache/spark/sql/execution/ShuffledRowRDD.scala

@@ -185,8 +185,6 @@ class ShuffledRowRDD(
       case CoalescedPartitionSpec(startReducerIndex, endReducerIndex) =>
         SparkEnv.get.shuffleManager.getReader(
           dependency.shuffleHandle,
-          0,
-          Int.MaxValue,
           startReducerIndex,
           endReducerIndex,
           context,