如何使用 PySpark 用组的聚合填充空值

Question

我有一个包含一些空值的简单数据集：

Age,Title
10,Mr
20,Mr
null,Mr
1, Miss
2, Miss
null, Miss

我想用不同列（在本例中为标题）的分组聚合填充空值。例如。标题列的平均值是：

15, Mr
1.5, Miss

所以最终的结果应该是这样的：

Age,Title
10,Mr
20,Mr
15,Mr
1, Miss
2, Miss
1.5, Miss

我已经看到很多使用 Pandas 使用 Transform 的示例：

df["Age"] = df.groupby("Title").transform(lambda x: x.fillna(x.mean()))

我试图不使用外部库，而是在 pyspark 中本地进行。 python数据框没有transform方法。

我正在考虑将聚合存储在一个单独的数据框中，如下所示：

meanAgeDf = df.groupBy("Title").mean("Age").select("Title", col("avg(Age)").alias("AgeMean"))

然后为每个分组查找标题并用该平均值填充所有这些值：

from pyspark.sql.functions import when, col
x = df.join(meanAgeDf, "Title").withColumn("AgeMean", when(col("Age").isNull(), col("AgeMean")).otherwise(col("Age")))

这是最有效的方法吗？

Answer 1

这可以使用窗口函数avg一步完成。

from pyspark.sql import Window
from pyspark.sql.functions import when,avg
w = Window.partitionBy(df.title)
res = df.withColumn("mean_col",avg(df.age).over(w))