如何按分组均值填充PySpark DataFrame列的null值

方案1：基于已计算好的均值表实现 #

假设你有两个DataFrame：

• 原始数据表：raw_df，包含highway、speed_kph两个字段
• 预先计算的均值表：avg_speed_df，包含highway、avg_speed两个字段，其中avg_speed就是你已经算好的各分组speed_kph均值（Road对应65、Service对应30）

实现代码如下：

from pyspark.sql.functions import coalesce

result_df = raw_df.join(avg_speed_df, on="highway", how="left") \
    # coalesce会返回第一个非空值：原有speed_kph非空就保留原值，为空就用对应分组均值
    .withColumn("speed_kph", coalesce("speed_kph", "avg_speed")) \
    # 删除多余的均值辅助字段
    .drop("avg_speed")

方案2：无需提前计算均值表，用窗口函数一步实现 #

如果你还没做分组聚合计算均值，可以直接用窗口函数完成全流程，代码更简洁：

from pyspark.sql.functions import avg, coalesce
from pyspark.sql.window import Window

# 定义按highway字段分组的窗口
window_spec = Window.partitionBy("highway")

result_df = raw_df.withColumn("speed_kph", 
                coalesce("speed_kph", avg("speed_kph").over(window_spec))
              )

两种方案最终输出的结果都和你给出的预期结果完全一致。

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