如何在PySpark透视表中按列计算百分比（优先聚合阶段实现）

背景数据 #

原始CSV数据如下：

desc,B1,B2,B3,B4,B5,B6,B7,B8,B9,B10,B11
Other,111957.0,35293.0,225852.0,35110.0,1023680.0,448736.0,256473.0,269856.0,306668.0,8807.0,89551.0
Down,575614.0,203186.0,0.0,125056.0,0.0,766086.0,1157311.0,11127.0,88741.0,31603.0,300733.0
Up,0.0,0.0,1953645.0,0.0,346423.0,0.0,0.0,0.0,0.0,0.0,0.0
Same,2948065.0,730113.0,33121.0,668868.0,5451224.0,4485121.0,30780025.0,1977361.0,5295598.0,217697.0,1790024.0
Old,186596.0,88257.0,0.0,36842.0,2173626.0,240619.0,0.0,2770.0,2212560.0,9865.0,121045.0
New,0.0,0.0,0.0,0.0,3148.0,0.0,97252.0,0.0,0.0,0.0,0.0

该DataFrame由以下代码生成：

y = x.groupby('desc').pivot('prev_segment').sum('cust_count')

需求 #

将每个B*列的数值转换为该列总和的百分比，使每列总计为100%，优先在透视的聚合函数中完成。

解决方案 #

通过窗口函数预先计算每个prev_segment的全局总客户数，再计算单条记录占该分组的百分比，最后透视聚合百分比即可。代码如下：

from pyspark.sql import Window
import pyspark.sql.functions as F

# 定义窗口：按prev_segment分区，计算该分区的总客户数
segment_window = Window.partitionBy('prev_segment')

# 1. 计算每个prev_segment对应的总客户数
x_with_total = x.withColumn('segment_total', F.sum('cust_count').over(segment_window))

# 2. 计算当前记录占该segment的百分比，保留2位小数
x_with_pct = x_with_total.withColumn('cust_pct', F.round(F.col('cust_count') / F.col('segment_total') * 100, 2))

# 3. 替换原sum聚合，直接透视聚合百分比
y_pct = x_with_pct.groupby('desc').pivot('prev_segment').sum('cust_pct')

说明 #

• 窗口函数partitionBy('prev_segment')会对每个prev_segment（即最终的B*列）单独计算总和，确保百分比计算的准确性。
• 用round函数控制百分比的小数位数，避免精度问题；若不需要保留小数，可去掉该函数。
• 最终生成的y_pct中，每列的数值为对应desc在该prev_segment中的占比，列总和约为100%（因四舍五入可能存在±0.01的误差）。

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