Spark Shuffle的类型分类及核心实现详解

// Hash Shuffle 的简化实现逻辑
class HashShuffleWriter[K, V] {
  private val numReducers = dep.partitioner.numPartitions
  
  def write(records: Iterator[Product2[K, V]]): Unit = {
    val writers = Array.fill(numReducers)(new DiskBlockObjectWriter)
    
    records.foreach { record =>
      val partitionId = dep.partitioner.getPartition(record._1)
      writers(partitionId).write(record)
    }
    
    writers.foreach(_.close())
  }
}

// Sort Shuffle 的核心流程
class SortShuffleWriter[K, V, C] {
  private val sorter = new ExternalSorter[K, V, C](
    dep.aggregator,
    dep.keyOrdering,
    dep.serializer
  )
  
  def write(records: Iterator[Product2[K, V]]): Unit = {
    // 将数据插入外部排序器
    sorter.insertAll(records)
    
    // 获取输出文件
    val outputFile = shuffleBlockResolver.getDataFile(dep.shuffleId, mapId)
    val tmp = Utils.tempFileWith(outputFile)
    
    try {
      val blockId = ShuffleBlockId(dep.shuffleId, mapId, IndexShuffleBlockResolver.NOOP_REDUCE_ID)
      
      // 写入数据和索引文件
      val partitionLengths = sorter.writePartitionedFile(blockId, tmp)
      shuffleBlockResolver.writeIndexFileAndCommit(dep.shuffleId, mapId, partitionLengths, tmp)
    } finally {
      if (tmp.exists() && !tmp.delete()) {
        logError(s"Error while deleting temp file ${tmp.getAbsolutePath}")
      }
    }
  }
}

// Tungsten Shuffle 的内存管理
class TungstenShuffleWriter[K, V] {
  private val memoryManager = new TaskMemoryManager(
    new ExecutorMemoryManager(new SparkConf())
  )
  
  private val sorter = new UnsafeExternalSorter(
    memoryManager,
    blockManager,
    serializerManager,
    taskContext,
    recordComparator,
    prefixComparator,
    initialSize,
    pageSizeBytes,
    numPartitionsForShuffle
  )
  
  def write(records: Iterator[Product2[K, V]]): Unit = {
    records.foreach { record =>
      val partitionId = partitioner.getPartition(record._1)
      sorter.insertRecord(
        record._1,
        record._2,
        partitionId,
        false // don't need to copy data
      )
    }
    
    // 获取排序后的数据迭代器
    val sortedIterator = sorter.getSortedIterator
    
    // 写入输出文件
    val outputFile = shuffleBlockResolver.getDataFile(shuffleId, mapId)
    val writer = new DiskBlockObjectWriter(outputFile, serializerManager, 
      bufferSize, syncWrites, writeMetrics)
    
    try {
      sortedIterator.foreach { record =>
        writer.write(record)
      }
    } finally {
      writer.close()
    }
  }
}

// 优化前：产生大量 Shuffle 数据
val result = rdd.groupByKey()
  .mapValues(_.sum)
 
// 优化后：使用 reduceByKey 减少数据传输
val result = rdd.reduceByKey(_ + _)

// 根据数据量调整分区数
val optimalPartitions = Math.max(
  200, // 最小分区数
  (totalDataSize / 128MB).toInt // 每个分区约 128MB
)
 
val optimizedRDD = originalRDD
  .coalesce(optimalPartitions, shuffle = true)
  .mapPartitions(processPartition)

// 当小表与大表 join 时，使用广播变量
val smallTable = spark.read.parquet("small_table.parquet")
val broadcastSmallTable = spark.sparkContext.broadcast(
  smallTable.collect().toMap
)
 
val result = largeTable.mapPartitions { partition =>
  val smallMap = broadcastSmallTable.value
  partition.map { record =>
    // 使用广播的小表数据进行 join
    (record.key, record.value, smallMap.get(record.key))
  }
}

# Spark 配置文件中的 Shuffle 优化参数
spark.sql.shuffle.partitions=200
spark.shuffle.file.buffer=64KB
spark.shuffle.spill.compress=true
spark.shuffle.compress=true
spark.shuffle.service.enabled=true
spark.shuffle.registration.timeout=120s
spark.shuffle.registration.maxAttempts=3

// 自定义 Shuffle 监控器
class ShufflePerformanceMonitor extends SparkListener {
  override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = {
    val shuffleMetrics = taskEnd.taskMetrics.shuffleReadMetrics
    
    // 监控 Shuffle 读取数据量
    val shuffleBytesRead = shuffleMetrics.remoteBytesRead
    val shuffleRecordsRead = shuffleMetrics.recordsRead
    
    // 监控 Shuffle 写入数据量
    val shuffleWriteMetrics = taskEnd.taskMetrics.shuffleWriteMetrics
    val shuffleBytesWritten = shuffleWriteMetrics.bytesWritten
    
    // 记录性能指标
    logInfo(s"Task ${taskEnd.taskInfo.taskId} - " +
      s"Shuffle Read: ${shuffleBytesRead} bytes, " +
      s"Shuffle Write: ${shuffleBytesWritten} bytes")
    
    // 性能告警
    if (shuffleBytesRead > 1e9) { // 超过 1GB
      logWarning(s"Large shuffle read detected: ${shuffleBytesRead} bytes")
    }
  }
}

# 使用 spark-submit 提交作业时启用详细监控
spark-submit \
  --conf spark.eventLog.enabled=true \
  --conf spark.eventLog.dir=/tmp/spark-events \
  --conf spark.history.fs.logDirectory=/tmp/spark-events \
  --class com.example.SparkJob \
  my-spark-job.jar
 
# 分析 Shuffle 性能
spark-submit \
  --class org.apache.spark.tools.ShuffleAnalyzer \
  spark-tools.jar \
  --event-log /tmp/spark-events/application_xxx \
  --output /tmp/shuffle-analysis

// 1. 优先使用窄依赖操作
val optimized = rdd.map(x => (x._1, x._2))
  .filter(_._2 > 0)  // 窄依赖
  .reduceByKey(_ + _)  // 宽依赖，但必要
 
// 2. 合理使用缓存避免重复 Shuffle
val expensiveRDD = someExpensiveTransformation()
expensiveRDD.persist(StorageLevel.MEMORY_AND_DISK)
 
// 3. 使用 salting 技术解决数据倾斜
val saltedRDD = originalRDD.map { case (key, value) =>
  val salt = Random.nextInt(10)
  ((key, salt), value)
}
 
val result = saltedRDD
  .reduceByKey(aggregationFunction)
  .map { case ((key, _), value) => (key, value) }
  .reduceByKey(aggregationFunction)