Как заполнить нулевые значения агрегатом группы с помощью PySpark

Question

У меня есть простой набор данных с некоторыми значениями NULL:

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

Я хочу заполнить значения NULL агрегатом группировки по другому столбцу (в данном случае, Title).Например, среднее значение столбца заголовка:

15, Mr
1.5, Miss

Таким образом, конечный результат должен выглядеть следующим образом:

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

Я видел много примеров с использованием Панд используя Преобразование :

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))