Spark DataFrame循环Join后操作卡顿：400万行数据显示耗时超10分钟求助

你遇到的这个性能问题，核心原因是循环执行join+drop的操作让Spark的执行计划变得极度臃肿，冗余的Shuffle操作累积导致性能崩盘。

为什么循环操作会变慢？ #

每一次循环里的groupby().agg()都会触发一次Shuffle（把数据按分组键重新分区），紧接着的join又会再一次涉及数据的重分区或匹配操作。哪怕你最后删掉了新增的求和列，Spark的查询优化器也没办法完全消除这些冗余的Shuffle和关联逻辑——它会把每一步的操作都记录在执行计划里。当列数超过6个时，执行计划的复杂度会线性增长，最终哪怕是简单的display，Spark都要先解析这个臃肿的计划，再执行大量无意义的计算，自然慢到无法忍受。

解决方案1：一次性完成聚合与关联 #

最直接的优化方式是把所有列的聚合操作一次性完成，再只做一次join，这样只需要一次Shuffle，执行计划会简洁得多：

from pyspark.sql.functions import sum

# 一次性生成所有列的求和聚合表达式
agg_exprs = [sum(c).alias(f'sum_{c}') for c in list_columns]
# 一次分组聚合得到所有列的分组求和结果
sum_agg_df = df.groupby(list_group_features).agg(*agg_exprs)
# 仅执行一次join关联回原DataFrame
df = df.join(sum_agg_df, list_group_features)
# 一次性删除所有求和列
df = df.drop(*[f'sum_{c}' for c in list_columns])

这种方式把原来N次的Shuffle+join操作，压缩成1次，性能提升会非常明显。

解决方案2：使用窗口函数替代join #

如果你的场景允许，还可以用窗口函数来实现分组求和，完全避免join操作，效率会更高：

from pyspark.sql.functions import sum
from pyspark.sql.window import Window

# 定义分组窗口
group_window = Window.partitionBy(list_group_features)

# 循环计算分组求和后删除列（仅需一次Shuffle）
for c in list_columns:
    df = df.withColumn(f'sum_{c}', sum(c).over(group_window))
    df = df.drop(f'sum_{c}')

窗口函数会在一次Shuffle后完成所有分组内的计算，不需要再做join关联，执行计划会更简洁，性能也更优。

关键提醒 #

虽然最终的DataFrame和初始状态结构、数据一致，但Spark不会自动识别并消除这些冗余的中间操作——它只会严格按照你编写的步骤生成执行计划。所以要尽量避免这种“做了又删”的循环操作，改用批量处理或更高效的算子来替代。

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