如何优化PySpark中的GROUPBY操作以加速大数据集处理？

针对你处理数十亿行数据时的性能问题，以下是具体的优化方案：

1. 前置数据过滤（最核心优化） #

原代码在分组聚合后才过滤数据，导致Spark需要先处理所有历史数据再过滤，极大浪费计算资源。必须将过滤操作提前到分组之前，只保留需要的近600天数据，直接减少后续分组聚合的数据量：

# 先过滤，再处理后续逻辑
df_filtered = df.filter(col("Date") >= date_sub(current_date(), 600))

2. 移除冗余的select操作 #

原代码在聚合后执行了全字段select，但此时DataFrame已经包含所有需要的字段（分组字段+聚合字段+新增的Year/Month），这个select完全冗余，直接删除即可，减少不必要的算子开销。

3. 替换first聚合为分组字段（如果业务逻辑允许） #

观察你的聚合逻辑：对M-S字段使用first()，如果每个(Date,B,C,D,E,F,G,H,I,J,K,L)分组下的M-S值是唯一的，可以直接将M-S加入groupBy列表，彻底避免聚合操作——聚合是Spark中开销较高的算子，替换后性能会大幅提升：

df_grouped = (
    df_filtered
    .groupBy("Date", "B", "C", "D", "E", "F", "G", "H", "I", "J", "K", "L", "M", "N", "O", "P", "Q", "R", "S")
    .count()  # 仅用于完成groupBy，后续删除该列
    .drop("count")
    .withColumn("Year", year(col("Date")))
    .withColumn("Month", month(col("Date")))
)

如果业务上分组内M-S存在重复值，必须保留first()，可以尝试使用min()或max()替代——部分情况下这些聚合函数的执行效率比first()更高（取决于数据分布），同时可以添加ignoreNulls=True参数减少无效计算：

first("M", ignoreNulls=True).alias("M")

4. 利用分区裁剪（如果源表是分区表） #

如果你的源表是按Date、Year或Month分区的，确保过滤条件直接命中分区字段，让Spark只读取需要的分区数据，避免全表扫描：

# 假设源表按Year和Month分区，直接过滤分区字段
df_filtered = df.filter(
    (col("Year") >= year(date_sub(current_date(), 600))) &
    (col("Month") >= month(date_sub(current_date(), 600)))
)

5. 优化Spark并行度配置 #

针对数十亿行数据，默认的spark.sql.shuffle.partitions（200）会导致每个分区数据量过大，增加shuffle开销。根据集群CPU核数调整该参数，一般设置为核数的2-3倍：

# 在代码开头设置（或通过spark-submit参数传递）
spark.conf.set("spark.sql.shuffle.partitions", 1000)  # 示例值，根据集群规模调整

6. 提前裁剪不必要的列 #

如果源DataFrame包含除(Date,B,C,D,E,F,G,H,I,J,K,L,M,N,O,P,Q,R,S)之外的其他列，先执行select只保留需要的列，减少数据传输和处理的体积：

df_trimmed = df.select("Date", "B", "C", "D", "E", "F", "G", "H", "I", "J", "K", "L", "M", "N", "O", "P", "Q", "R", "S")
df_filtered = df_trimmed.filter(col("Date") >= date_sub(current_date(), 600))

7. 调整Year/Month的计算时机 #

可以将Year和Month的计算提前到过滤之后、分组之前，避免分组后再处理：

df_preprocessed = (
    df_filtered
    .withColumn("Year", year(col("Date")))
    .withColumn("Month", month(col("Date")))
)

优化后的完整代码示例 #

# 配置并行度
spark.conf.set("spark.sql.shuffle.partitions", 1000)

df_grouped = (
    # 1. 裁剪列+过滤数据
    df.select("Date", "B", "C", "D", "E", "F", "G", "H", "I", "J", "K", "L", "M", "N", "O", "P", "Q", "R", "S")
    .filter(col("Date") >= date_sub(current_date(), 600))
    # 2. 提前计算Year/Month
    .withColumn("Year", year(col("Date")))
    .withColumn("Month", month(col("Date")))
    # 3. 替换聚合为全字段分组（如果业务允许）
    .groupBy("Date", "B", "C", "D", "E", "F", "G", "H", "I", "J", "K", "L", "M", "N", "O", "P", "Q", "R", "S", "Year", "Month")
    .count()
    .drop("count")
    # 如果业务必须用first聚合，替换为下面的逻辑：
    # .groupBy("Date", "B", "C", "D", "E", "F", "G", "H", "I", "J", "K", "L")
    # .agg(
    #     first("M", ignoreNulls=True).alias("M"),
    #     first("N", ignoreNulls=True).alias("N"),
    #     first("O", ignoreNulls=True).alias("O"),
    #     first("P", ignoreNulls=True).alias("P"),
    #     first("Q", ignoreNulls=True).alias("Q"),
    #     first("R", ignoreNulls=True).alias("R"),
    #     first("S", ignoreNulls=True).alias("S")
    # )
)

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