PySpark优化：高效计算各分类列层级数据占比

原代码的核心问题 #

你的当前实现存在几个关键性能瓶颈：

• 全局count()触发额外全表扫描，且无法正确计算分组内的占比（原代码中sz是全局总数，不是每个box的数量）
• 循环处理每个分类列并多次join，会产生大量Shuffle操作，在100个分类列的场景下性能极差
• 每个列单独执行pivot再合并，重复计算逻辑且浪费资源

高效Spark化实现方案 #

以下方案通过窗口函数+长格式转换+一次性Pivot完成计算，仅需2次全表扫描，避免循环和多次Shuffle：

import pyspark.sql.functions as f
from pyspark.sql.window import Window

# 1. 读取原始数据
box_of_potatos = spark.read... # 替换为你的数据读取逻辑

# 2. 按box分组，计算每组的总条数（窗口函数，仅一次扫描）
box_window = Window.partitionBy("box")
df_with_total = box_of_potatos.withColumn(
    "group_total", f.count("*").over(box_window)
)

# 3. 获取所有分类列（排除box列）
cat_cols = [c for c, t in box_of_potatos.dtypes if t.startswith("string") and c != "box"]

# 4. 将多分类列转为长格式（stack函数一次性完成列转行）
stack_expr = f"stack({len(cat_cols)}, {', '.join([f'{repr(c)}, {c}' for c in cat_cols])}) as (col_name, col_value)"
melted_df = df_with_total.select("box", "group_total", f.expr(stack_expr))

# 5. 拼接目标列名（如potato.red）
melted_df = melted_df.withColumn(
    "pivot_col", f.concat(f.col("col_name"), f.lit("."), f.col("col_value"))
)

# 6. 分组聚合计算占比，一次性Pivot转宽格式
result_df = melted_df.groupBy("box").pivot("pivot_col").agg(
    f.round(f.sum(f.lit(1)/f.col("group_total")), 2).alias("proportion")
).fillna(0)

# 查看结果
result_df.show()

针对大规模数据的额外优化 #

如果你的分类列层级极多（100个列×100层级=10000个目标列），可以提前指定Pivot的所有可能值，避免Spark额外扫描数据获取Pivot列：

# 预获取所有可能的Pivot列名（仅需一次扫描所有分类列的去重值）
all_pivot_cols = []
for col in cat_cols:
    distinct_vals = [row[0] for row in box_of_potatos.select(col).distinct().collect()]
    all_pivot_cols += [f"{col}.{val}" for val in distinct_vals]

# 传入预定义的Pivot列列表，提升性能
result_df = melted_df.groupBy("box").pivot("pivot_col", all_pivot_cols).agg(
    f.round(f.sum(f.lit(1)/f.col("group_total")), 2).alias("proportion")
).fillna(0)

方案优势 #

1. 减少扫描次数：仅2次全表扫描（窗口计算+聚合Pivot），远少于原方案的多次扫描
2. 避免循环Shuffle：一次性处理所有分类列，仅一次Shuffle操作
3. 分组占比准确：用窗口函数计算每个box的组内总数，替代全局count
4. 代码简洁可维护：无需拆分多个函数，逻辑统一清晰

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