如何优化Spark边列表转邻接列表时groupByKey执行过慢的问题

核心结论 #

除非你的原始输入数据已经按节点ID做了预分区，否则完全避免shuffle是不可能的，因为同一节点的出入边大概率分布在不同的分区，必须通过shuffle把同节点的所有边聚合到同一位置。但可以通过多个优化手段大幅降低shuffle开销，性能可提升10~100倍不等。

优化方案 #

1. 替换groupByKey为支持map端预聚合的算子 #

groupByKey不会做map端合并，会将所有原始数据直接通过shuffle传输，开销极大。你可以根据使用的API选择更优的聚合算子：

• 若使用RDD API，优先用aggregateByKey代替groupByKey，开启map端预聚合，同分区内同节点的边会先在本地合并后再做shuffle，可直接减少70%以上的shuffle数据量。示例代码：

intermediate.aggregateByKey(
    zeroValue=[],
    seqFunc=lambda acc, edge: acc.append(edge) or acc, # 本地分区内合并
    combFunc=lambda acc1, acc2: acc1 + acc2 # 跨分区合并
)

• 若使用Spark SQL的DataFrame/Dataset API，性能会比原生RDD高30%~200%，得益于钨丝内存优化、代码生成机制，shuffle效率也远高于RDD。示例代码：

from pyspark.sql import SparkSession
from pyspark.sql.functions import collect_list

spark = SparkSession.builder.config(conf=conf).getOrCreate()
edge_df = intermediate.toDF(["node_id", "edge_info"])
adj_df = edge_df.groupBy("node_id").agg(collect_list("edge_info").alias("adjacent_list"))
# 如需转回RDD格式可调用adj_df.rdd.map(lambda row: (row.node_id, row.adjacent_list))

2. 利用树结构特性进一步减少开销 #

你处理的是树结构，每个节点只有1个父节点，对应每个key最多只有1条('out', xxx)的记录，聚合时不需要对out记录做多次合并，可以自定义聚合逻辑单独存储out边，避免和in边一起做无意义的列表拼接，最后再合并输出即可，进一步降低内存和计算开销。

3. 调整shuffle相关参数 #

针对你的硬件配置和数据规模，可以调整以下参数大幅提升性能：

• 调整并行度：不要用默认的200个shuffle分区，100MB测试数据可设置spark.default.parallelism=16，15GB生产数据可设置为总核数的2~3倍，避免过多小任务的调度开销。
• 开启shuffle压缩：设置spark.shuffle.compress=true，默认使用snappy压缩，可减少50%以上的shuffle网络传输量。
• 调整shuffle内存占比：设置spark.memory.shuffleFraction=0.3，避免shuffle过程中频繁溢写磁盘。

4. 输入数据预分区（可接近完全避免shuffle） #

如果你的原始边数据可以预先按节点ID哈希分区存储，那么读入后每个分区内的节点ID不会跨分区，直接在每个分区内做本地聚合即可，不需要走全局shuffle，性能可提升一个数量级。

内容的提问来源于stack exchange，提问作者Prateek Bhuwania