[SPARK][CORE] 面試問題之 Shuffle reader 的細枝末節 （上）

歡迎關注微信公眾號“Tim在路上”
之前我們已經了解了shuffle writer的詳細過程，那么生成文件后會發生什么呢？以及它們是如何被讀取呢？讀取是內存的操作嗎？這些問題也隨之產生，那么今天我們將先來了解了shuffle reader的細枝末節。

在文章Spark Shuffle概述中我們已經知道，在ShuffleManager中不僅定義了getWriter來獲取map writer的實現方式， 同時還定義了getReader來獲取讀取shuffle文件的實現方式。 在Spark中調用有兩個調用getReader的抽象類的重要實現，分別是ShuffledRDD和ShuffleRowRDD。前者是與RDD API交互，后面一個是DataSet Api的交互實現。在Spark 3.0后其核心已經變成了Spark SQL，所以我們重點從ShuffleRowRDD調用getReader開始講起。

從ShuffleRowRDD開始

ShuffleRowRDD主要是被ShuffleExchangeExec調用。這里簡單介紹下ShuffleExchangeExec操作算子。它主要負責兩件事：首先，準備ShuffleDependency，它根據父節點所需的分區方案對子節點的輸出行進行分區。其次，添加一個ShuffleRowRDD并指定準備好的ShuffleDependency作為此RDD的依賴項。


2927.png
class ShuffledRowRDD(
    var dependency: ShuffleDependency[Int, InternalRow, InternalRow],
    metrics: Map[String, SQLMetric],
    partitionSpecs: Array[ShufflePartitionSpec])
  extends RDD[InternalRow](dependency.rdd.context,Nil)

ShuffleRowRDD繼承自RDD[InternalRow], 同時內部維護著三個參數，分別是dependency，metrics和partitionSpecs。dependency封裝著shuffleId ，shuffleHandle ，numPartitions 可以基于其判斷出shuffleWriter采用了哪種方式。partitionSpecs定義了分區規范的類型。

目前在spark 3.2版本中partitionSpecs的實現類主要有以下四個：

  • CoalescedPartitionSpec用于coalesce shuffle partitions 邏輯規則
  • PartialReducerPartitionSpec參與了 skew join 優化
  • PartialMapperPartitionSpec用于本地隨機讀取器
  • CoalescedMapperPartitionSpec用于優化本地隨機讀取器

不同類型的分區規范其實質是代表不同的隨機讀取的參數。我們都知道在Spark Shuffle中getReader僅有且唯一的一個實現方式, 即BlockStoreShuffleReader 的實現。但是不同的分區規范意味將給共享的reader器傳遞不同的參數， 下面是ShuffleRowRDD中的簡化代碼：

// ShuffleRowRDD
override def compute(split: Partition, context: TaskContext): Iterator[InternalRow] = {
  val tempMetrics = context.taskMetrics().createTempShuffleReadMetrics()
  // `SQLShuffleReadMetricsReporter` will update its own metrics for SQL exchange operator,
  // as well as the `tempMetrics` for basic shuffle metrics.
  val sqlMetricsReporter = new SQLShuffleReadMetricsReporter(tempMetrics, metrics)
  val reader = split.asInstanceOf[ShuffledRowRDDPartition].spec match {
    // CoalescedPartitionSpec會讀取map task為所有reducer所產生的shuffle file
    case CoalescedPartitionSpec(startReducerIndex, endReducerIndex, _) =>
      SparkEnv.get.shuffleManager.getReader(
        dependency.shuffleHandle,
        startReducerIndex,
        endReducerIndex,
        context,
        sqlMetricsReporter)
   // PartialReducerPartitionSpec 讀取map task為一個reducer產生的部分數據
    case PartialReducerPartitionSpec(reducerIndex, startMapIndex, endMapIndex, _) =>
      SparkEnv.get.shuffleManager.getReader(
        dependency.shuffleHandle,
        startMapIndex,
        endMapIndex,
        reducerIndex,
        reducerIndex + 1,
        context,
        sqlMetricsReporter)
   // PartialMapperPartitionSpec讀取shuffle map文件的部分
   case PartialMapperPartitionSpec(mapIndex, startReducerIndex, endReducerIndex) =>
        SparkEnv.get.shuffleManager.getReader(
          dependency.shuffleHandle,
          mapIndex,
          mapIndex + 1,
          startReducerIndex,
          endReducerIndex,
          context,
          sqlMetricsReporter)
...
    reader.read().asInstanceOf[Iterator[Product2[Int, InternalRow]]].map(_._2)
  }

其實從上面傳的參數中就可以看出點端倪CoalescedPartitionSpec(startReducerIndex,endReducer-Index) 讀取map task為所有reducer所產生的shuffle file；PartialReducerPartitionSpec(startMap-Index, endMapIndex,reducerIndex,reducerIndex + 1) 可以看出每次讀取只會為一個reducer讀取部分數據。

從上面代碼可以看出ShuffleRowRDD 使用 read() 方法遍歷 shuffle 數據并將其返回給客戶端，那么接下來我們就詳細的看下getReader是如何實現的？

ShuffleReader調用前的準備

SortShuffleManager是ShuffleManager的唯一實現，里面也實現getReader方法，那么就讓我們從getReader開始。

override def getReader[K, C](
    handle: ShuffleHandle,
    startMapIndex: Int,
    endMapIndex: Int,
    startPartition: Int,
    endPartition: Int,
    context: TaskContext,
    metrics: ShuffleReadMetricsReporter): ShuffleReader[K, C] = {
  val baseShuffleHandle = handle.asInstanceOf[BaseShuffleHandle[K, _, C]]
  val (blocksByAddress, canEnableBatchFetch) =
    // 是否開啟了push-based shuffle, 后續再分享，這里先跳過
    if (baseShuffleHandle.dependency.shuffleMergeEnabled) {
      val res = SparkEnv.get.mapOutputTracker.getPushBasedShuffleMapSizesByExecutorId(
        handle.shuffleId, startMapIndex, endMapIndex, startPartition, endPartition)
      (res.iter, res.enableBatchFetch)
    } else {
      // [1] 使用mapOutputTracker獲取shuffle塊的位置
      val address = SparkEnv.get.mapOutputTracker.getMapSizesByExecutorId(
        handle.shuffleId, startMapIndex, endMapIndex, startPartition, endPartition)
      (address, true)
    }
  // [2] 創建一個BlockStoreShuffleReader實例，該實例將負責將shuffle文件從mapper傳遞到 reducer 任務
  new BlockStoreShuffleReader(
    handle.asInstanceOf[BaseShuffleHandle[K, _, C]], blocksByAddress, context, metrics,
    shouldBatchFetch =
      canEnableBatchFetch &&canUseBatchFetch(startPartition, endPartition, context))
}

可以看到getReader主要做了兩件事：

  • [1] 使用mapOutputTracker獲取shuffle塊的位置
  • [2] 創建一個BlockStoreShuffleReader實例，該實例將負責將shuffle文件從mapper傳遞到reducer 任務

那么Spark中如何保存和獲取shuffle塊的位置呢？

在spark中有兩種mapOutputTracker，兩種mapOutputTracker 都是在創建SparkEnv時創建。

其中第一個是MapOutputTrackerMaster，它駐留在驅動程序中并跟蹤每個階段的map output輸出, 并與DAGScheduler進行通信。

另一個是MapOutputTrackerWorker，位于執行器上，它負責從MapOutputTrackerMaster獲取shuffle 元數據信息。

MapOutputTrackerMaster：

  1. DAGScheduler在創建 shuffle map 階段時會調用registerShuffle方法，從下面的代碼可以看出在創建ShuffleMapStage會調用registerShuffle，其實質是在向 shuffleStatuses 映射器中放入shuffleid, 并為其值創建一個新的new ShuffleStatus(numMaps)。
def createShuffleMapStage[K, V, C](
    shuffleDep: ShuffleDependency[K, V, C], jobId: Int): ShuffleMapStage = {
  val rdd = shuffleDep.rdd
  ...
  stageIdToStage(id) = stage
  shuffleIdToMapStage(shuffleDep.shuffleId) = stage
  updateJobIdStageIdMaps(jobId, stage)

  if (!mapOutputTracker.containsShuffle(shuffleDep.shuffleId)) {
    // 在創建ShuffleMapStage會調用registerShuffle
    mapOutputTracker.registerShuffle(shuffleDep.shuffleId, rdd.partitions.length,
      shuffleDep.partitioner.numPartitions)
  }
  stage
}

def registerShuffle(shuffleId: Int, numMaps: Int, numReduces: Int): Unit = {
    if (pushBasedShuffleEnabled) {
      if (shuffleStatuses.put(shuffleId, new ShuffleStatus(numMaps, numReduces)).isDefined) {
        throw new IllegalArgumentException("Shuffle ID " + shuffleId + " registered twice")
      }
    } else {
      // 可以看到其實質是在向 shuffleStatuses 放入shuffleid, 創建ShuffleStatus
      if (shuffleStatuses.put(shuffleId, new ShuffleStatus(numMaps)).isDefined) {
        throw new IllegalArgumentException("Shuffle ID " + shuffleId + " registered twice")
      }
    }
  }
  1. 到目前位置master tracker存放了一個shuffleid, 表明DAG中存在一個shuffle, 但還是不知道map output file的具體位置。
// DAGScheduler中
private[scheduler] def handleTaskCompletion(event: CompletionEvent): Unit = {

  case smt: ShuffleMapTask =>
     val shuffleStage = stage.asInstanceOf[ShuffleMapStage]
     ...
     mapOutputTracker.registerMapOutput(
        shuffleStage.shuffleDep.shuffleId, smt.partitionId, status)
  }

def registerMapOutput(shuffleId: Int, mapIndex: Int, status: MapStatus): Unit = {
    shuffleStatuses(shuffleId).addMapOutput(mapIndex, status)
}

從上面代碼可以看出，在每次 shuffle map 階段的任務終止時，DAGScheduler都會向MapOutputTrackerMaster發送狀態更新。跟蹤器將有關特定 shuffle 文件的位置和大小的信息添加到在注冊步驟中初始化 的shuffleStatuses map中。


3tled.png

MapOutputTrackerWorker：

當worker tracker 沒有緩存shuffle信息, 這時就必須發送GetMapOutputStatuses消息來從master tracker獲取它。

再回過頭來看看，在getReader中通過mapOutputTracker獲取shuffle塊的位置的方法。

// mapOutTracker
private def getMapSizesByExecutorIdImpl(
    shuffleId: Int,
    startMapIndex: Int,
    endMapIndex: Int,
    startPartition: Int,
    endPartition: Int,
    useMergeResult: Boolean): MapSizesByExecutorId = {
  logDebug(s"Fetching outputs for shuffle$shuffleId")
  // [1] 獲取mapOutputStatuses
  val (mapOutputStatuses, mergedOutputStatuses) = getStatuses(shuffleId, conf,
    // EnableBatchFetch can be set to false during stage retry when the
    // shuffleDependency.shuffleMergeEnabled is set to false, and Driver
    // has already collected the mergedStatus for its shuffle dependency.
    // In this case, boolean check helps to insure that the unnecessary
    // mergeStatus won't be fetched, thus mergedOutputStatuses won't be
    // passed to convertMapStatuses. See details in [SPARK-37023].
    if (useMergeResult)fetchMergeResultelse false)
  ...
}

從上面可以看出獲取具體的map output 位置的實現在getStatuses方法中。下面我們來具體分析下：

private def getStatuses(
    shuffleId: Int,
    conf: SparkConf,
    canFetchMergeResult: Boolean): (Array[MapStatus], Array[MergeStatus]) = {
  // push-based shuffle 開啟，獲取MergeStatus， 現暫不考慮
  if (canFetchMergeResult) {
    ...
  } else {
    val statuses = mapStatuses.get(shuffleId).orNull
    // [1] 如果mapStatuses不包含statuses， 就向master tracker發送GetMapOutputStatuses消息
    if (statuses == null) {
      logInfo("Don't have map outputs for shuffle " + shuffleId + ", fetching them")
      val startTimeNs = System.nanoTime()
fetchingLock.withLock(shuffleId) {
        var fetchedStatuses =mapStatuses.get(shuffleId).orNull
        if (fetchedStatuses == null) {
          logInfo("Doing the fetch; tracker endpoint = " +trackerEndpoint)
          val fetchedBytes = askTracker[Array[Byte]](GetMapOutputStatuses(shuffleId))
          try {
            fetchedStatuses =
              MapOutputTracker.deserializeOutputStatuses[MapStatus](fetchedBytes, conf)
          } catch {
            ...
          }
          logInfo("Got the map output locations")
          mapStatuses.put(shuffleId, fetchedStatuses)
        }
        (fetchedStatuses, null)
      }
    // [2] 如果mapStatuses包含statuses, 直接返回
    } else {
      (statuses, null)
    }
  }
}

從getStatuses可以看出：

  • [1] 如果mapStatuses不包含statuses, 就向master tracker發送GetMapOutputStatuses消息
  • [2] 如果mapStatuses包含statuses, 直接返回
private[spark] sealed trait MapStatus extends ShuffleOutputStatus {
  def location: BlockManagerId

  def updateLocation(newLoc: BlockManagerId): Unit

  def getSizeForBlock(reduceId: Int): Long

  def mapId: Long
}

可見MapStatus中包含了location， mapId等信息。

最后，回到getReader方法中，通過SparkEnv.get.mapOutputTracker.getMapSizesByExecutorId獲取shuffle塊信息后，再將其作為 shuffle 塊的及其物理位置傳遞給BlockStoreShuffleReader。

那么接下來就我們再來分析下BlockStoreShuffleReader的實現。

為避免冗長將BlockStoreShuffleReader放到下一講進行分析。
歡迎關注微信公眾號“Tim在路上”

?著作權歸作者所有,轉載或內容合作請聯系作者
平臺聲明：文章內容（如有圖片或視頻亦包括在內）由作者上傳并發布，文章內容僅代表作者本人觀點，簡書系信息發布平臺，僅提供信息存儲服務。

推薦閱讀更多精彩內容