上一篇解决了 DAGScheduler 和 TaskScheduler 的分工。这一篇进入 Shuffle 机制——Stage 之间数据交换的物理过程。

Shuffle 容易被理解成"把数据从一组节点传到另一组节点"。更准确的说法是:Shuffle 是一个分布式的排序-分区-传输流水线,它把上游 Stage 每个分区的输出按 key 的目标分区号排序写入磁盘,然后由下游 Stage 的 Task 跨网络拉取属于自己的那部分数据。

本文只抓一个问题:SortShuffleManager 的 Shuffle Write 和 Shuffle Read 两个阶段各做了什么,以及为什么 Shuffle 是 Spark 作业的头号性能瓶颈。

Shuffle 的全局视角

一个 Shuffle 操作(如 reduceByKey)在物理层面涉及两组 Task:

1
2
3
4
5
6
7
8
9
10
11
12
13
上游 Stage (M 个 ShuffleMapTask)          下游 Stage (R 个 Task)
┌──────────┐ ┌──────────┐
│ Map T-0 │──┬──────────────────────────>│ Red T-0 │
│ │ │ ┌──>│ │
├──────────┤ │ ┌───────────────────┤ ├──────────┤
│ Map T-1 │──┼──┤ │ │ Red T-1 │
│ │ │ │ ┌──>│ │ │
├──────────┤ │ │ ┌────────────┤ │ ├──────────┤
│ Map T-2 │──┘ └──┤ │ └──>│ Red T-2 │
│ │ └────────────┘ │ │
└──────────┘ └──────────┘
Shuffle Write Shuffle Read
(排序 + 写磁盘) (拉取 + 归并)

每个 Map Task 把输出数据按目标分区号分成 R 份。每个 Reduce Task 从所有 M 个 Map Task 那里拉取属于自己的那份数据。整个 Shuffle 涉及 M×R 次数据传输。

Shuffle Write:排序与落盘

Spark 当前默认的 Shuffle 实现是 SortShuffleManager。Shuffle Write 阶段的流程:

Map Task 计算出一条记录后,根据 key 的 hash 值计算目标分区号,然后把 (partitionId, key, value) 写入一个内存缓冲区(PartitionedAppendOnlyMap 或 PartitionedPairBuffer)。

当内存缓冲区达到阈值时,按 (partitionId, key) 排序后溢写(spill)到磁盘文件。一个 Map Task 可能产生多次溢写,生成多个临时文件。

Map Task 结束时,把所有溢写文件和内存中剩余的数据做一次归并排序(merge sort),产生一个最终的 Shuffle 数据文件(.data)和一个索引文件(.index)。索引文件记录了每个目标分区在数据文件中的偏移量和长度。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
Map Task 的 Shuffle Write 流程:

记录 → 计算 partitionId → 写入内存缓冲区

缓冲区满? ──┤
│ 是 │ 否
▼ │
排序 + 溢写到磁盘 │
│ │
▼ │
Task 结束 ◄───────┘


归并所有溢写文件 + 内存数据


输出:data 文件 + index 文件

关键细节:每个 Map Task 只产生一个 data 文件和一个 index 文件,而不是 R 个文件。这是 Sort-Based Shuffle 相对早期 Hash-Based Shuffle 的核心改进——Hash-Based Shuffle 每个 Map Task 产生 R 个文件,总共 M×R 个文件,在大规模作业中(M=10000, R=10000)会产生 1 亿个小文件,压垮文件系统。

Shuffle Read:拉取与归并

下游 Stage 的 Task 启动后,首先从 Driver 的 MapOutputTracker 查询所有 Map Task 的 Shuffle 输出位置(哪个 Executor、哪个文件、哪个偏移量)。然后向对应的 Executor 发起 HTTP 请求,拉取属于自己分区的数据。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
Reduce Task 的 Shuffle Read 流程:

查询 MapOutputTracker
│ 获取所有 Map 输出位置

并行拉取数据(默认 5 个并发连接)
│ 从各 Executor 的 Shuffle 文件中
│ 按 index 文件定位偏移量

内存缓冲 + 溢写(如果数据量大)


归并排序(如果需要排序)
│ 或直接聚合(如果是 reduceByKey)

输出迭代器给后续算子

Shuffle Read 的数据拉取是并行的,默认同时从 5 个 Map Task 拉取(spark.reducer.maxSizeInFlight 控制缓冲区大小,spark.reducer.maxReqsInFlight 控制并发请求数)。拉取到的数据先放入内存缓冲区,如果内存不够则溢写到磁盘。

对于 reduceByKey 这样的聚合操作,Shuffle Read 不需要全局排序——只需要按 key 聚合。Spark 使用 ExternalAppendOnlyMap 进行哈希聚合,当内存不足时溢写到磁盘。

Map 端预聚合

reduceByKey 和 groupByKey 都触发 Shuffle,但两者的 Shuffle 数据量可能差异巨大。区别在于 reduceByKey 在 Map 端做了预聚合(combine)。

1
2
3
4
5
6
7
8
9
10
数据: [("a",1), ("b",1), ("a",1), ("a",1), ("b",1)]

groupByKey (无预聚合):
Shuffle Write: ("a",1), ("a",1), ("a",1), ("b",1), ("b",1) → 5 条记录过网络
Shuffle Read: ("a", [1,1,1]), ("b", [1,1])

reduceByKey (有预聚合):
Map 端先聚合: ("a",3), ("b",2)
Shuffle Write: ("a",3), ("b",2) → 2 条记录过网络
Shuffle Read: ("a",3), ("b",2) (如果只有一个 Map Task)

Map 端预聚合可以大幅减少 Shuffle 传输的数据量,尤其是在 key 的基数(distinct 值数量)远小于记录总数时。这也是 Spark 社区推荐用 reduceByKey 替代 groupByKey 的原因。

Shuffle 的三重代价

磁盘 I/O:Shuffle Write 必须把数据写入本地磁盘(溢写文件和最终 data 文件),Shuffle Read 拉取后如果内存不足也要溢写。一个大作业可能在每个节点上产生 GB 级别的 Shuffle 临时文件。

网络传输:Shuffle Read 需要跨网络拉取数据。在一个 1000 节点的集群上,如果 Shuffle 数据量是 1TB,理论上需要在集群内部传输 1TB 数据。网络带宽和延迟成为瓶颈。

序列化开销:数据在写入磁盘和网络传输前需要序列化,读取后需要反序列化。Spark 默认使用 Java 序列化(效率较低),切换到 Kryo 序列化可以显著减少 Shuffle 数据大小和 CPU 消耗。

Shuffle 的演进

Spark 的 Shuffle 实现经历了三代演进:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
Hash Shuffle (Spark 0.x - 1.1)
每个 Map Task 为每个 Reduce 分区创建一个文件
总文件数: M × R
问题: 大规模作业产生海量小文件

Sort Shuffle (Spark 1.1+, 默认)
每个 Map Task 输出一个排序后的 data 文件 + index 文件
总文件数: 2 × M
改进: 文件数从 M×R 降到 2×M

Tungsten Sort Shuffle (Spark 1.4+, 自动触发)
使用堆外内存和二进制排序
避免 Java 对象反序列化
条件: 无 map 端聚合 + 分区数 < 16777216 + 序列化器支持 relocation

当前 Spark 版本中,SortShuffleManager 是唯一的 Shuffle 管理器,但内部会根据条件自动选择 UnsafeShuffleWriter(Tungsten 路径)或 SortShuffleWriter(普通路径)。

实验:观察 Shuffle 指标

1
2
3
4
5
6
7
8
9
10
// 准备测试数据
val data = sc.parallelize(1 to 1000000).map(i => (i % 100, i))

// 使用 groupByKey(无 map 端预聚合)
val grouped = data.groupByKey(10)
grouped.count()

// 使用 reduceByKey(有 map 端预聚合)
val reduced = data.reduceByKey(_ + _, 10)
reduced.count()

在 Spark UI 的 Stages 页面对比两个作业的 Shuffle Write 数据量。reduceByKey 的 Shuffle Write 应该远小于 groupByKey,因为 map 端预聚合把每个分区的 100 个 key 先聚合成了少量记录。

进入 Stage 详情页,查看 Summary Metrics 中的 Shuffle Write Size / Records 和 Shuffle Read Size / Records,可以直观看到预聚合的效果。

模式提炼

Shuffle 体现了分布式数据重分布的通用模式:先本地排序/分桶,再跨网络传输,最后目标端归并。

这个模式不只出现在 Spark 中。MapReduce 的 Shuffle 做同样的事情:Map 输出按 key 排序写入本地磁盘,Reduce Task 从所有 Map Task 拉取并归并。数据库的分布式 Hash Join 也遵循相同的流程:两表按 join key 重分布到相同的节点上,然后本地 join。

优化 Shuffle 的通用策略也是跨系统通用的:减少数据量(预聚合、过滤前置)、减少文件数(合并小文件)、减少序列化开销(二进制格式)、避免不必要的 Shuffle(co-partition 优化)。

工程迁移表

Spark 概念 MapReduce 对应 Flink 对应 数据库对应
Shuffle Write Map 输出写本地磁盘 上游算子写 Network Buffer Hash 分布写入临时表
Shuffle Read Reduce 拉取 Map 输出 下游算子从 Buffer 读取 远端节点拉取分布式表数据
SortShuffleManager MapReduce 的 Sort-based Shuffle Flink 的 Hash/Sort Shuffle 数据库的 Exchange 算子实现
Map 端预聚合(Combine) Combiner 接口 无直接对应(算子自行优化) Partial Aggregation
External Shuffle Service 无(MapReduce 不分离 Shuffle 服务) 无(Flink Shuffle 在 TaskManager 内) 独立的 Shuffle 服务(如 Remote Shuffle Service)

常见误解

误解一:“Shuffle 数据通过 Driver 中转”。Shuffle 数据在 Executor 之间直接传输(点对点),不经过 Driver。Driver 只负责提供 Map 输出的位置信息(MapOutputTracker),不传输实际数据。

误解二:“增加分区数总能提升 Shuffle 性能”。增加下游分区数(reduce 端并行度)可以减小每个 Reduce Task 处理的数据量,但同时增加了每个 Map Task 需要分桶的数量和 Shuffle Read 的连接数。存在一个平衡点。

误解三:“Shuffle 数据可以缓存在内存中避免磁盘 I/O”。Shuffle Write 的输出必须写入磁盘——这是 Spark 容错机制的要求。如果 Executor 故障,内存中的 Shuffle 数据会丢失,但磁盘上的 Shuffle 文件可以被其他 Executor 读取(通过 External Shuffle Service)。

练习

  1. 对同一份数据分别调用 groupByKey 和 reduceByKey,在 Spark UI 中对比 Shuffle Write 的 Records 和 Bytes 指标。计算预聚合带来的数据压缩比。

  2. 修改 spark.sql.shuffle.partitions(默认 200)为 10 和 1000,对同一个 SparkSQL 查询执行三次。在 Spark UI 中观察 Shuffle Read 的 Task 数量和总执行时间的变化。

  3. 配置 spark.serializer=org.apache.spark.serializer.KryoSerializer,对同一个 Shuffle 作业执行两次(默认 Java 序列化 vs Kryo),对比 Shuffle Write 的数据大小。

系列导航

序号 主题 状态
00 导读:为什么 Spark 的核心是一张 DAG 已完成
01 RDD:弹性分布式数据集的五大属性 已完成
02 宽依赖与窄依赖:Stage 是怎样划分出来的 已完成
03 DAG Scheduler 与 Task Scheduler:从逻辑计划到物理执行 已完成
04 Shuffle 机制:数据跨分区交换的代价与优化 本篇
05 内存管理与 Tungsten:堆外内存、序列化与代码生成 下一篇
06 存储体系:Block Manager、广播变量与累加器
07 容错机制:Lineage、Checkpoint 与推测执行
08 Spark SQL 与 Catalyst:从 SQL 文本到物理算子树
09 DataFrame 与 Dataset:类型安全与性能的折中
10 Adaptive Query Execution:运行时改写执行计划
11 数据源 API 与谓词下推:让存储层帮忙过滤
12 Structured Streaming:把流当成无界表
13 水位线与状态管理:流处理的时间语义
14 Spark on YARN / K8s:资源管理与部署模式
15 动态资源分配与作业调度:多租户集群的资源博弈
16 性能诊断:Spark UI、Event Log 与调优思路
17 Spark 与 Flink:两种计算引擎的设计选择

参考资料

  • Apache Spark 源码:core/src/main/scala/org/apache/spark/shuffle/sort/SortShuffleManager.scala https://github.com/apache/spark
  • Apache Spark 源码:core/src/main/scala/org/apache/spark/shuffle/sort/SortShuffleWriter.scala
  • Databricks Blog: Optimizing Apache Spark’s Shuffle Operation. https://www.databricks.com/blog
  • Matei Zaharia. An Architecture for Fast and General Data Processing on Large Clusters. PhD Thesis, UC Berkeley, 2014. Chapter 5.