PySpark：循环生成多周统计列并合并DataFrame的解决方法

需要给列表weeks里的每个周数生成对应的统计列，然后把所有列整合到同一个DataFrame里。目前用循环处理每个周，但每次循环都会覆盖之前的结果，最后只保留了最后一个周的数据，想通过左连接把所有周的列合并到一起。

用户尝试的代码 #

weeks = [24, 25]
for weeknum in weeks:
    df_new = df.filter(df.week == weeknum).groupBy(['gender', 'pro']).pivot("share").agg(first('forecast_units')) \
        .withColumnRenamed('0.01', 'units_1_share_wk'+str(weeknum))
        .withColumnRenamed('0.1', 'units_10_share_wk'+str(weeknum))
        .withColumnRenamed('0.15', 'units_15_share_wk'+str(weeknum))
        .withColumnRenamed('0.2', 'units_20_share_wk'+str(weeknum)) 
df_new.show()

原DataFrame结构 #

|country|gender|order_date|         pro|share|        prediction|week|dayofweek|forecast_units|
+-------+------+----------+------------+------+------------------+----+---------+-------------------+
| ES|  Male|2022-09-15|Jeans - Flat| 0.01|13.322306632995605|  37|        5|               93.0|
| ES|  Male|2022-09-15|Jeans - Flat|  0.1| 19.09369468688965|  37|        5|              134.0|
| ES|  Male|2022-09-15|Jeans - Flat| 0.15|22.504554748535156|  37|        5|              158.0|

期望的DataFrame结构 #

|gender|pro|units_1_share_wk24|units_10_share_wk24|units_15_share_wk24|units_20_share_wk24|units_1_share_wk25|units_10_share_wk25|units_15_share_wk25|units_20_share_wk25|

期望输出示例 #

|gender|pro|units_1_share_wk24|units_10_share_wk24|units_15_share_wk24|units_20_share_wk24|units_1_share_wk25|units_10_share_wk25|units_15_share_wk25|units_20_share_wk25|
|---+---+---+---+---+---+---+---+---+---+
|Female|Belts|28.0|0.0|0.0|0.0|28.0|0.0|0.0|0.0|
|Female|Dress|0.0|44.0|0.0|0.0|0.0|0.0|0.0|0.0|
|Male|Belts|0.0|0.0|33.0|0.0|28.0|0.0|0.0|0.0|
|Male|Suits|0.0|0.0|0.0|34.0|0.0|0.0|0.0|0.0|

方法一：循环左连接整合 #

先拿第一个周的数据作为基础，之后每个周处理完都和基础DataFrame做左连接，把新列加进去，避免覆盖之前的结果：

from pyspark.sql import functions as F

weeks = [24, 25]
# 初始化结果DF，取第一个周的数据
result_df = df.filter(df.week == weeks[0]).groupBy(['gender', 'pro']).pivot("share").agg(F.first('forecast_units')) \
    .withColumnRenamed('0.01', 'units_1_share_wk24') \
    .withColumnRenamed('0.1', 'units_10_share_wk24') \
    .withColumnRenamed('0.15', 'units_15_share_wk24') \
    .withColumnRenamed('0.2', 'units_20_share_wk24')

# 循环处理剩下的周，逐个左连接
for weeknum in weeks[1:]:
    temp_df = df.filter(df.week == weeknum).groupBy(['gender', 'pro']).pivot("share").agg(F.first('forecast_units')) \
        .withColumnRenamed('0.01', 'units_1_share_wk'+str(weeknum)) \
        .withColumnRenamed('0.1', 'units_10_share_wk'+str(weeknum)) \
        .withColumnRenamed('0.15', 'units_15_share_wk'+str(weeknum)) \
        .withColumnRenamed('0.2', 'units_20_share_wk'+str(weeknum))
    # 按gender和pro做外连接，保证所有组合都保留
    result_df = result_df.join(temp_df, on=['gender', 'pro'], how='outer')

# 把空值填成0，和期望输出一致
result_df = result_df.fillna(0)
result_df.show()

方法二：一次性多维度pivot（更高效） #

不用循环连接，直接按gender、pro分组，同时pivotweek和share的组合，批量重命名列，性能比循环连接好很多：

from pyspark.sql import functions as F

weeks = [24, 25]
# 先过滤出目标周的数据
filtered_df = df.filter(df.week.isin(weeks))

# 按gender、pro分组，pivot week和share的组合
result_df = filtered_df.groupBy(['gender', 'pro']) \
    .pivot(['week', 'share']) \
    .agg(F.first('forecast_units'))

# 批量重命名列，转换成想要的格式
def rename_col(col_name):
    # 拆分pivot生成的列名，比如"(24, 0.01)"
    week_num, share_val = col_name.strip('()').split(', ')
    # 映射share值到对应的命名
    share_map = {'0.01': '1', '0.1': '10', '0.15': '15', '0.2': '20'}
    return f'units_{share_map[share_val]}_share_wk{week_num}'

# 重命名所有非分组列
new_columns = [rename_col(col) if col not in ['gender', 'pro'] else col for col in result_df.columns]
result_df = result_df.toDF(*new_columns)

# 填充空值为0
result_df = result_df.fillna(0)
result_df.show()

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