PySpark循环未生成DataFrame问题排查与解决

初始DataFrame df 结构如下：

+-------+---+-----+------------------+----+-------------------+
|gender| pro|share|        prediction|week|     forecast_units|
+------+----+-----+------------------+----+-------------------+
|  Male|Polo| 0.01| 258.4054260253906|  37|             1809.0|
|  Male|Polo|  0.1| 332.4026794433594|  38|             2327.0|
|  Male|Polo| 0.15|425.97430419921875|  39|             2982.0|
|  Male|Polo|  0.2| 508.3385314941406|  40|             3558.0|
....

尝试通过以下代码从原DataFrame生成多个DataFrame：先创建指定Schema的四个空DataFrame并存入list_dfs；再循环遍历周数列表c_weeks（37、38、39、40），对每个周数执行过滤、分组、透视聚合操作并赋值给list_dfs中的元素，但最终所有DataFrame仍为空。请问如何修复？若此方法不可行，该如何实现需求？

代码如下：

schema = StructType([\
    StructField("gender", StringType(),True), \
    StructField("pro",StringType(),True), \
    StructField("units_1_tpr",DoubleType(),True), \
    StructField("units_1'_tpr",DoubleType(),True), \
    StructField("units_15_tpr",DoubleType(),True), \
    StructField("units_20_tpr",DoubleType(),True)])

df_wk1 = spark.createDataFrame([],schema=schema)
df_wk2 = spark.createDataFrame([],schema=schema)
df_wk3 = spark.createDataFrame([],schema=schema)
df_wk4 = spark.createDataFrame([],schema=schema)

list_dfs = [df_wk1, df_wk2, df_wk3, df_wk4]
c_weeks = [37, 38, 39, 40]

for data,weeknum in zip(list_dfs, campaign_weeks):
    data = df.filter(df.week == weeknum).groupBy(['gender', 'pro']).pivot("share").agg(first('forecast_units'))

问题根源 #

循环里的data = ...只是给循环变量重新赋值，并没有修改list_dfs里的原始空DataFrame对象。Spark DataFrame是不可变对象，赋值操作不会改变列表中原本的空DataFrame，所以最终列表里还是初始的空数据。

修复方法 #

直接生成处理后的DataFrame并替换列表元素，不需要预先创建空DataFrame，具体有两种实现方式：

方式1：修改循环逻辑 #

list_dfs = []
c_weeks = [37, 38, 39, 40]

for weeknum in c_weeks:
    # 过滤+分组透视
    processed_df = df.filter(df.week == weeknum).groupBy(['gender', 'pro']).pivot("share").agg(first('forecast_units'))
    # 重命名列以匹配目标Schema
    processed_df = processed_df.withColumnRenamed("0.01", "units_1_tpr")\
                               .withColumnRenamed("0.1", "units_1'_tpr")\
                               .withColumnRenamed("0.15", "units_15_tpr")\
                               .withColumnRenamed("0.2", "units_20_tpr")
    list_dfs.append(processed_df)

# 拆分到单独变量（可选）
df_wk1, df_wk2, df_wk3, df_wk4 = list_dfs

方式2：列表推导式（更简洁） #

c_weeks = [37, 38, 39, 40]

def process_single_week(weeknum):
    temp_df = df.filter(df.week == weeknum).groupBy(['gender', 'pro']).pivot("share").agg(first('forecast_units'))
    # 重命名列
    return temp_df.withColumnRenamed("0.01", "units_1_tpr")\
                  .withColumnRenamed("0.1", "units_1'_tpr")\
                  .withColumnRenamed("0.15", "units_15_tpr")\
                  .withColumnRenamed("0.2", "units_20_tpr")

# 生成结果列表
list_dfs = [process_single_week(wk) for wk in c_weeks]
df_wk1, df_wk2, df_wk3, df_wk4 = list_dfs

补充说明 #

• 透视后的列名默认是share字段的原始值（0.01、0.1等），必须通过withColumnRenamed修改为目标Schema中的列名，否则列名不匹配
• 不需要预先创建空DataFrame，直接生成处理后的DataFrame存入列表即可，避免不可变对象赋值的问题
• 如果某周的share值不全，透视后对应列会显示null，符合目标Schema中字段允许为空的设置

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