如何判断PySpark中需重分区的时机以规避数据倾斜？

一、如何确定重分区的时机避免数据倾斜 #

重分区的核心是平衡节点数据负载，重点关注以下阶段：

• 宽依赖操作前后：join、groupBy、distinct这类操作是倾斜重灾区。若发现join的key分布不均、groupBy的聚合key存在热点值，提前对大表重分区可打散倾斜数据；如果操作后已出现倾斜（比如部分任务执行时间远超其他任务），也需重分区修复。
• 数据量骤变节点：经过过滤、聚合后数据量大幅增减时，原分区数不再适配（比如过滤后剩余数据仅为原来的10%，但分区数不变会产生大量小分区），此时重分区能平衡负载。
• 最终输出前：写入HDFS、S3等存储前重分区，可避免生成大量小文件或超大文件，同时保证输出阶段任务负载均匀。

二、通过Web UI检查Executor状态与任务数 #

完全可以通过Spark Web UI查看相关指标：

• Executors页面（默认端口4040）：可查看每个Executor的内存使用率（已用/总内存）、完成任务数、失败任务数、GC情况等，直观判断节点负载压力。
• Stages页面：能查看每个Stage的总任务数，以及单个任务的输入数据量、执行时间。若大部分任务几秒完成，个别任务耗时几十分钟，基本可判定存在数据倾斜。

三、编程式检查分区倾斜的优化方案 #

你提到的代码存在明显问题：collect()会触发全量Action，数据倾斜时会拖慢速度甚至导致Driver OOM。可以通过以下方式优化：

1. 抽样统计分区大小 #

只抽取部分分区数据估算整体分布，大幅降低开销：

# 统计每个分区的记录数，返回RDD[Int]
partition_sizes_rdd = df.rdd.mapPartitions(lambda iter: [len(list(iter))])

# 抽样10%的分区数据（可根据实际调整抽样比例）
sampled_sizes = partition_sizes_rdd.sample(withReplacement=False, fraction=0.1).collect()

if sampled_sizes:
    max_size = max(sampled_sizes)
    min_size = min(sampled_sizes)
    # 设定倾斜阈值（比如最大分区是最小的5倍以上则判定为倾斜）
    if max_size / min_size > 5:
        # 根据目标分区大小（比如128MB/分区）计算合理的分区数
        total_records = df.count()
        avg_record_size = df.rdd.map(lambda x: len(str(x))).mean()  # 估算单条记录大小
        target_partitions = int((total_records * avg_record_size) / (128 * 1024 * 1024)) + 1
        df = df.repartition(target_partitions)

2. 利用近似统计API #

Spark的stat模块提供近似统计方法，无需全量计算：

# 先获取各分区大小的RDD
partition_sizes_rdd = df.rdd.mapPartitions(lambda iter: [len(list(iter))])

# 将RDD转为DataFrame以便使用stat API
from pyspark.sql import Row
size_df = spark.createDataFrame(partition_sizes_rdd.map(lambda x: Row(size=x)))

# 获取分区大小的近似分位数（比如95%分位数和5%分位数）
quantiles = size_df.stat.approxQuantile("size", [0.05, 0.95], 0.05)
lower_bound = quantiles[0]
upper_bound = quantiles[1]

# 如果95%分位数是5%分位数的数倍以上，判定为倾斜
if upper_bound / lower_bound > 5:
    df = df.repartition(int(df.rdd.count() / (128 * 1024 * 1024)) + 1)

额外建议：事前预防倾斜 #

比起事后修复，事前预防更高效：

• 对倾斜的key做加盐处理（给热点key拼接随机后缀，打散后再聚合）；
• 小表使用广播join（broadcast(df)），避免大表 shuffle；
• 拆分热点key，单独处理后再合并结果。

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