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

Question

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

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

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

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数据帧没有变换方法。

我想将聚合存储在一个单独的数据帧中，如下所示：

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 import when,avg
w = Window.partitionBy(df.title)
res = df.withColumn("mean_col",avg(df.age).over(w))