[SPARK-32922][SHUFFLE][CORE] Adds support for executors to fetch local and remote merged shuffle data

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

@@ -1453,7 +1453,8 @@ private[spark] object MapOutputTracker extends Logging {
             // ShuffleBlockId with mapId being SHUFFLE_PUSH_MAP_ID to indicate this is
             // a merged shuffle block.
             splitsByAddress.getOrElseUpdate(mergeStatus.location, ListBuffer()) +=
-              ((ShuffleBlockId(shuffleId, SHUFFLE_PUSH_MAP_ID, partId), mergeStatus.totalSize, -1))
+              ((ShuffleBlockId(shuffleId, SHUFFLE_PUSH_MAP_ID, partId), mergeStatus.totalSize,
+                SHUFFLE_PUSH_MAP_ID))
             // For the "holes" in this pre-merged shuffle partition, i.e., unmerged mapper
             // shuffle partition blocks, fetch the original map produced shuffle partition blocks
             val mapStatusesWithIndex = mapStatuses.zipWithIndex

core/src/main/scala/org/apache/spark/serializer/SerializerManager.scala

@@ -110,6 +110,7 @@ private[spark] class SerializerManager(
   private def shouldCompress(blockId: BlockId): Boolean = {
     blockId match {
       case _: ShuffleBlockId => compressShuffle
+      case _: ShuffleBlockChunkId => compressShuffle
       case _: BroadcastBlockId => compressBroadcast
       case _: RDDBlockId => compressRdds
       case _: TempLocalBlockId => compressShuffleSpill

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

@@ -35,6 +35,7 @@ private[spark] class BlockStoreShuffleReader[K, C](
     readMetrics: ShuffleReadMetricsReporter,
     serializerManager: SerializerManager = SparkEnv.get.serializerManager,
     blockManager: BlockManager = SparkEnv.get.blockManager,
+    mapOutputTracker: MapOutputTracker = SparkEnv.get.mapOutputTracker,
     shouldBatchFetch: Boolean = false)
   extends ShuffleReader[K, C] with Logging {
 
@@ -71,6 +72,7 @@ private[spark] class BlockStoreShuffleReader[K, C](
       context,
       blockManager.blockStoreClient,
       blockManager,
+      mapOutputTracker,
       blocksByAddress,
       serializerManager.wrapStream,
       // Note: we use getSizeAsMb when no suffix is provided for backwards compatibility

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

@@ -43,6 +43,7 @@ sealed abstract class BlockId {
     (isInstanceOf[ShuffleBlockId] || isInstanceOf[ShuffleBlockBatchId] ||
      isInstanceOf[ShuffleDataBlockId] || isInstanceOf[ShuffleIndexBlockId])
   }
+  def isShuffleChunk: Boolean = isInstanceOf[ShuffleBlockChunkId]
   def isBroadcast: Boolean = isInstanceOf[BroadcastBlockId]
 
   override def toString: String = name
@@ -72,6 +73,15 @@ case class ShuffleBlockBatchId(
   }
 }
 
+@Since("3.2.0")
+@DeveloperApi
+case class ShuffleBlockChunkId(
+    shuffleId: Int,
+    reduceId: Int,
+    chunkId: Int) extends BlockId {
+  override def name: String = "shuffleChunk_" + shuffleId  + "_" + reduceId + "_" + chunkId
+}
+
 @DeveloperApi
 case class ShuffleDataBlockId(shuffleId: Int, mapId: Long, reduceId: Int) extends BlockId {
   override def name: String = "shuffle_" + shuffleId + "_" + mapId + "_" + reduceId + ".data"
@@ -152,14 +162,15 @@ 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]+)".r
   val SHUFFLE_BATCH = "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 SHUFFLE_PUSH = "shufflePush_([0-9]+)_([0-9]+)_([0-9]+)".r
   val SHUFFLE_MERGED_DATA = "shuffleMerged_([_A-Za-z0-9]*)_([0-9]+)_([0-9]+).data".r
   val SHUFFLE_MERGED_INDEX = "shuffleMerged_([_A-Za-z0-9]*)_([0-9]+)_([0-9]+).index".r
   val SHUFFLE_MERGED_META = "shuffleMerged_([_A-Za-z0-9]*)_([0-9]+)_([0-9]+).meta".r
+  val SHUFFLE_CHUNK = "shuffleChunk_([0-9]+)_([0-9]+)_([0-9]+)".r
   val BROADCAST = "broadcast_([0-9]+)([_A-Za-z0-9]*)".r
   val TASKRESULT = "taskresult_([0-9]+)".r
   val STREAM = "input-([0-9]+)-([0-9]+)".r
@@ -186,6 +197,8 @@ object BlockId {
       ShuffleMergedIndexBlockId(appId, shuffleId.toInt, reduceId.toInt)
     case SHUFFLE_MERGED_META(appId, shuffleId, reduceId) =>
       ShuffleMergedMetaBlockId(appId, shuffleId.toInt, reduceId.toInt)
+    case SHUFFLE_CHUNK(shuffleId, reduceId, chunkId) =>
+      ShuffleBlockChunkId(shuffleId.toInt, reduceId.toInt, chunkId.toInt)
     case BROADCAST(broadcastId, field) =>
       BroadcastBlockId(broadcastId.toLong, field.stripPrefix("_"))
     case TASKRESULT(taskId) =>

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

@@ -132,6 +132,11 @@ private[spark] class HostLocalDirManager(
       executorIdToLocalDirsCache.asMap().asScala.toMap
     }
 
+  private[spark] def getCachedHostLocalDirsFor(executorId: String): Option[Array[String]] =
+    executorIdToLocalDirsCache.synchronized {
+      Option(executorIdToLocalDirsCache.getIfPresent(executorId))
+    }
+
   private[spark] def getHostLocalDirs(
       host: String,
       port: Int,