PySpark：验证同列重分区的DataFrame连接是否触发Shuffle

环境信息 #

• Spark版本：2.4.7
• 已设置spark.conf.set("spark.sql.autoBroadcastJoinThreshold", -1)

背景 #

为避免Shuffle，我将两个PySpark DataFrame按相同列、相同分区数重分区后写入Parquet文件，现需验证后续连接操作是否仍会触发Shuffle。

初始DataFrame统计 #

>>> initital_df1.count(), initial_df1.select('id').distinct().count()  # 每个id对应唯一行
(1000, 1000)
>>> initial_df2.count(), initial_df2.select('id').distinct().count()  # 每个id对应300行（25年月度数据）
(300000, 1000)
>>> initial_df1.rdd.getNumPartitions(), initial_df2.rdd.getNumPartitions()
(2, 8)
>>> initial_df1.rdd.glom().map(len).collect()
[500, 500]
>>> initial_df2.rdd.glom().map(len).collect()
[56400, 50700, 45000, 40800, 36300, 30000, 25200, 15600]

操作步骤 #

步骤1：按相同列和分区数重分区并写入 #

>>> initial_df1.repartition(4, 'id').write.parquet(df1_write_path)
>>> initial_df2.repartition(4, 'id').write.parquet(df2_write_path)

步骤2：读取重分区后的DataFrame #

>>> repartitioned_df1 = spark.read.parquet(df1_write_path)
>>> repartitioned_df2 = spark.read.parquet(df2_write_path)

重分区后DataFrame统计 #

>>> repartitioned_df1.rdd.getNumPartitions(), repartitioned_df2.rdd.getNumPartitions()
(4, 4)
>>> initial_df1.rdd.glom().map(len).collect()
[279, 255, 238, 228]
>>> initial_df2.rdd.glom().map(len).collect()
[83700, 76500, 71400, 68400]

步骤3：连接重分区后的DataFrame并聚合写入 #

>>> agg_df = repartitioned_df1.join(repartitioned_df2, 'id').groupBy('id').agg(F.sum(col1_df1), F.avg(col1_df2))
>>> agg_df.write.parquet(agg_df_write_path)

测试场景 #

• 场景1：使用重分区后的repartitioned_df1和repartitioned_df2执行步骤3
• 场景2：使用初始initial_df1和initial_df2执行步骤3

观察结果 #

• 两种场景下agg_df.explain()生成的Spark执行计划一致
• SparkUI中Shuffle读写数据值一致
• 场景1的执行速度至少是场景2的2倍

问题 #

如何确认连接操作是再次触发哈希分区产生Shuffle，还是利用了预分区的共分区特性，仅移动写入阶段已创建的分区？

验证方法 #

1. 检查DataFrame的分区器信息

Spark RDD的分区器决定数据分布，按id重分区的DataFrame底层RDD应带有一致的HashPartitioner。执行以下代码查看：

# 查看重分区后DataFrame的分区器
print(repartitioned_df1.rdd.partitioner)
print(repartitioned_df2.rdd.partitioner)

若两者输出均为HashPartitioner(4)，说明使用了相同哈希分区策略，连接时可利用共分区特性避免Shuffle。

2. 分析物理执行计划的Exchange节点

执行agg_df.explain("extended")查看物理执行计划细节：

• 若连接操作（如SortMergeJoin/HashJoin）前无额外Exchange节点（Shuffle的标志），说明直接使用了预分区数据，未触发新Shuffle。

3. 对比Shuffle阶段任务数量

在SparkUI的Stages页面：

• 场景1中，若未触发新Shuffle，Shuffle阶段任务数应等于预分区数（4个）；
• 场景2中，初始DataFrame分区不一致，连接会触发Shuffle，任务数通常对应spark.sql.shuffle.partitions默认值（200）。

4. 检查任务本地性

共分区DataFrame连接时，同一id数据在同一Executor节点，任务本地性更高（如PROCESS_LOCAL）。在SparkUI的Tasks页面查看本地性统计：

• 场景1的PROCESS_LOCAL占比会远高于场景2，因无需跨节点传输数据。

5. 对比网络传输与磁盘IO

在SparkUI的Storage和Environment页面：

• 场景1的网络传输数据量远低于场景2，且无额外Shuffle临时文件生成；
• 场景2会因Shuffle产生大量网络传输和磁盘临时文件写入。

结论 #

若满足以下所有条件，说明连接操作利用了预分区的共分区特性，未触发新Shuffle：

• 两个DataFrame的分区器一致；
• 物理执行计划中连接前无额外Exchange节点；
• Shuffle任务数等于预分区数；
• 任务本地性占比高、网络传输量低。

内容的提问来源于stack exchange，提问作者Shivam Verma