[SPARK-9853][Core] Optimize shuffle fetch of contiguous partition IDs

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

@@ -812,10 +812,13 @@ private[spark] object MapOutputTracker extends Logging {
         logError(errorMessage)
         throw new MetadataFetchFailedException(shuffleId, startPartition, errorMessage)
       } else {
+        var totalSize = 0L
         for (part <- startPartition until endPartition) {
-          splitsByAddress.getOrElseUpdate(status.location, ArrayBuffer()) +=
-            ((ShuffleBlockId(shuffleId, mapId, part), status.getSizeForBlock(part)))
+          totalSize += status.getSizeForBlock(part)
         }
+        splitsByAddress.getOrElseUpdate(status.location, ArrayBuffer()) +=
+          ((ShuffleBlockId(shuffleId, mapId, startPartition, endPartition - startPartition),
+            totalSize))
       }
     }
 

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

@@ -196,12 +196,14 @@ private[spark] class IndexShuffleBlockResolver(
   override def getBlockData(blockId: ShuffleBlockId): ManagedBuffer = {
     // The block is actually going to be a range of a single map output file for this map, so
     // find out the consolidated file, then the offset within that from our index
+    logDebug(s"Fetch block data for $blockId")
     val indexFile = getIndexFile(blockId.shuffleId, blockId.mapId)
 
     val in = new DataInputStream(Files.newInputStream(indexFile.toPath))
     try {
       ByteStreams.skipFully(in, blockId.reduceId * 8)
       val offset = in.readLong()
+      ByteStreams.skipFully(in, (blockId.length - 1) * 8)
       val nextOffset = in.readLong()
       new FileSegmentManagedBuffer(
         transportConf,

core/src/main/scala/org/apache/spark/storage/BlockId.scala

@@ -52,8 +52,9 @@ case class RDDBlockId(rddId: Int, splitIndex: Int) extends BlockId {
 // Format of the shuffle block ids (including data and index) should be kept in sync with
 // org.apache.spark.network.shuffle.ExternalShuffleBlockResolver#getBlockData().
 @DeveloperApi
-case class ShuffleBlockId(shuffleId: Int, mapId: Int, reduceId: Int) extends BlockId {
-  override def name: String = "shuffle_" + shuffleId + "_" + mapId + "_" + reduceId
+case class ShuffleBlockId(shuffleId: Int, mapId: Int, reduceId: Int, length: Int = 1)
+  extends BlockId {
+  override def name: String = "shuffle_" + shuffleId + "_" + mapId + "_" + reduceId + "_" + length
 }
 
 @DeveloperApi
@@ -103,7 +104,7 @@ class UnrecognizedBlockId(name: String)
 @DeveloperApi
 object BlockId {
   val RDD = "rdd_([0-9]+)_([0-9]+)".r
-  val SHUFFLE = "shuffle_([0-9]+)_([0-9]+)_([0-9]+)".r
+  val SHUFFLE = "shuffle_([0-9]+)_([0-9]+)_([0-9]+)_([0-9]+)".r
   val SHUFFLE_DATA = "shuffle_([0-9]+)_([0-9]+)_([0-9]+).data".r
   val SHUFFLE_INDEX = "shuffle_([0-9]+)_([0-9]+)_([0-9]+).index".r
   val BROADCAST = "broadcast_([0-9]+)([_A-Za-z0-9]*)".r
@@ -116,8 +117,8 @@ object BlockId {
   def apply(name: String): BlockId = name match {
     case RDD(rddId, splitIndex) =>
       RDDBlockId(rddId.toInt, splitIndex.toInt)
-    case SHUFFLE(shuffleId, mapId, reduceId) =>
-      ShuffleBlockId(shuffleId.toInt, mapId.toInt, reduceId.toInt)
+    case SHUFFLE(shuffleId, mapId, reduceId, n) =>
+      ShuffleBlockId(shuffleId.toInt, mapId.toInt, reduceId.toInt, n.toInt)
     case SHUFFLE_DATA(shuffleId, mapId, reduceId) =>
       ShuffleDataBlockId(shuffleId.toInt, mapId.toInt, reduceId.toInt)
     case SHUFFLE_INDEX(shuffleId, mapId, reduceId) =>

core/src/main/scala/org/apache/spark/storage/ShuffleBlockFetcherIterator.scala

@@ -515,7 +515,7 @@ final class ShuffleBlockFetcherIterator(
 
   private def throwFetchFailedException(blockId: BlockId, address: BlockManagerId, e: Throwable) = {
     blockId match {
-      case ShuffleBlockId(shufId, mapId, reduceId) =>
+      case ShuffleBlockId(shufId, mapId, reduceId, n) =>
         throw new FetchFailedException(address, shufId.toInt, mapId.toInt, reduceId, e)
       case _ =>
         throw new SparkException(

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

@@ -531,7 +531,7 @@ class CleanerTester(
 
   private def getShuffleBlocks(shuffleId: Int): Seq[BlockId] = {
     blockManager.master.getMatchingBlockIds( _ match {
-      case ShuffleBlockId(`shuffleId`, _, _) => true
+      case ShuffleBlockId(`shuffleId`, _, _, _) => true
       case ShuffleIndexBlockId(`shuffleId`, _, _) => true
       case _ => false
     }, askSlaves = true)

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

@@ -275,4 +275,36 @@ class MapOutputTrackerSuite extends SparkFunSuite {
     }
   }
 
+  test("fetch contiguous partitions") {
+    val rpcEnv = createRpcEnv("test")
+    val tracker = newTrackerMaster()
+    tracker.trackerEndpoint = rpcEnv.setupEndpoint(MapOutputTracker.ENDPOINT_NAME,
+      new MapOutputTrackerMasterEndpoint(rpcEnv, tracker, conf))
+    tracker.registerShuffle(10, 2)
+    assert(tracker.containsShuffle(10))
+    val size1000 = MapStatus.decompressSize(MapStatus.compressSize(1000L))
+    val size2000 = MapStatus.decompressSize(MapStatus.compressSize(2000L))
+    val size10000 = MapStatus.decompressSize(MapStatus.compressSize(10000L))
+    tracker.registerMapOutput(10, 0, MapStatus(BlockManagerId("a", "hostA", 1000),
+      Array(1000L, 10000L, 2000L)))
+    tracker.registerMapOutput(10, 1, MapStatus(BlockManagerId("b", "hostB", 1000),
+      Array(10000L, 2000L, 1000L)))
+    val statuses1 = tracker.getMapSizesByExecutorId(10, 0, 2)
+    assert(statuses1.toSet ===
+      Seq((BlockManagerId("a", "hostA", 1000),
+        ArrayBuffer((ShuffleBlockId(10, 0, 0, 2), size1000 + size10000))),
+        (BlockManagerId("b", "hostB", 1000),
+          ArrayBuffer((ShuffleBlockId(10, 1, 0, 2), size10000 + size2000))))
+        .toSet)
+    val statuses2 = tracker.getMapSizesByExecutorId(10, 2, 3)
+    assert(statuses2.toSet ===
+      Seq((BlockManagerId("a", "hostA", 1000),
+        ArrayBuffer((ShuffleBlockId(10, 0, 2, 1), size2000))),
+        (BlockManagerId("b", "hostB", 1000),
+          ArrayBuffer((ShuffleBlockId(10, 1, 2, 1), size1000))))
+        .toSet)
+    assert(0 == tracker.getNumCachedSerializedBroadcast)
+    tracker.stop()
+    rpcEnv.shutdown()
+  }
 }

core/src/test/scala/org/apache/spark/storage/BlockIdSuite.scala

@@ -55,7 +55,7 @@ class BlockIdSuite extends SparkFunSuite {
     val id = ShuffleBlockId(1, 2, 3)
     assertSame(id, ShuffleBlockId(1, 2, 3))
     assertDifferent(id, ShuffleBlockId(3, 2, 3))
-    assert(id.name === "shuffle_1_2_3")
+    assert(id.name === "shuffle_1_2_3_1")
     assert(id.asRDDId === None)
     assert(id.shuffleId === 1)
     assert(id.mapId === 2)