PySpark按不同列多次执行groupBy运算效率低的优化方案咨询

你当前的循环逐列聚合写法性能差的核心原因是：每一列的groupBy都会触发一次全表扫描+一次shuffle，500列就会重复扫描大表500次、触发500次shuffle，IO和计算冗余度极高，大表场景下速度会非常慢。
下面给的优化方案可以将计算耗时降低1-2个数量级，同时完全适配你后续PSI计算、多维度（SEGMENT/Period）聚合的需求。 #

最优实现方案：列转行后单次聚合 #

核心逻辑是把宽表转成长表，只做1次全表扫描、1次shuffle就完成所有列的聚合计算，结果和你原有逻辑完全一致。

from pyspark.sql import functions as F

# 分组键配置，后续算Period维度时直接替换/追加字段即可，比如改成['SEGMENT', 'Period']
group_cols = ['SEGMENT']
# 自动过滤掉分组键，拿到所有需要做计数统计的业务列
calc_cols = [col for col in df_ref_crt.columns if col not in group_cols]

# 一步完成宽表转长表：把500个业务列统一转成「变量名+变量值」的行结构
df_long = df_ref_crt.select(
    *group_cols,
    F.explode(
        F.array(*[
            F.struct(
                F.lit(col_name).alias('variable'),
                F.col(col_name).cast("string").alias('new_name')
            ) for col_name in calc_cols
        ])
    ).alias('col_tmp')
).select(
    *group_cols,
    'col_tmp.variable',
    'col_tmp.new_name'
)

# 单次groupBy完成所有列、所有分组的计数，输出结果等价于你原代码中list_grp所有DataFrame union后的结果
df_agg = df_long.groupBy(*group_cols, 'variable', 'new_name').count()

说明：把所有变量值统一转为string是为了兼容不同列的数据类型（数值、字符串、枚举等），PSI计算本身只依赖取值分箱的计数结果，转string不会影响计算准确性，还能避免类型不兼容报错。 #

进一步提速的配套优化 #

• 缓存复用中间结果：如果你后续需要分别按SEGMENT、Period维度做聚合，转完长表df_long之后直接调用df_long.persist()做缓存，等所有聚合计算完成后再调用unpersist()释放缓存，避免重复做列转行计算。
• 提前裁剪数据：如果你的表是按Period等字段分区的，聚合前先filter出需要计算的时间周期数据，不要扫描全量历史分区。
• 开启Databricks原生自适应优化：计算前执行以下参数配置，自动处理shuffle倾斜、优化分区数：

SET spark.sql.adaptive.enabled = true;
SET spark.sql.adaptive.skewJoin.enabled = true;
SET spark.sql.aggregate.pushdown.enabled = true;

原写法性能损耗点说明 #

1. 每轮循环单独触发一次全表扫描，500列的磁盘IO开销是优化后方案的数百倍
2. 每轮循环触发一次独立shuffle，而shuffle是Spark计算中开销最高的环节，500次shuffle的延迟远高于单次shuffle
3. 原逻辑得到的是包含500个小DataFrame的列表，后续做PSI计算时还需要额外做union、关联等操作，会进一步增加计算开销，优化后直接输出结构化的聚合结果，可直接用于后续PSI占比计算、基准期对比等步骤。

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