замените NA медианой в pyspark, используя оконную функцию

Question

Я хочу заменить NA на medain, основываясь на столбцах разделов, используя оконную функцию в pyspark?

Пример ввода:

Требуемый выход:

Answer 1

Создание вашего фрейма данных:

list=([1,5,4],
    [1,5,None],
    [1,5,1],
    [1,5,4],
    [2,5,1],
    [2,5,2],
    [2,5,None],
    [2,5,None],
     [2,5,4])
df=spark.createDataFrame(list,['I_id','p_id','xyz'])
df.show()

+----+----+----+
|I_id|p_id| xyz|
+----+----+----+
|   1|   5|   4|
|   1|   5|null|
|   1|   5|   1|
|   1|   5|   4|
|   2|   5|   1|
|   2|   5|   2|
|   2|   5|null|
|   2|   5|null|
|   2|   5|   4|
+----+----+----+

Чтобы сохранить решение как можно более обобщенным c и динамическим c, мне пришлось создать много новых столбцов для вычисления медианы и возможности отправить его на нули. С учетом сказанного решение не будет медленным и будет масштабируемым для больших данных.

from pyspark.sql.window import Window
from pyspark.sql import functions as F
from pyspark.sql.functions import when
w= Window().partitionBy("I_id","p_id").orderBy(F.col("xyz").asc_nulls_first())
w2= Window().partitionBy("I_id","p_id")
df.withColumn("xyz1",F.count(F.col("xyz").isNotNull()).over(w))\
.withColumn("xyz2", F.max(F.row_number().over(w)).over(w2))\
.withColumn("xyz3", F.first("xyz1").over(w))\
.withColumn("xyz10", F.col("xyz2")-F.col("xyz3"))\
.withColumn("xyz9", F.when((F.col("xyz2")-F.col("xyz3"))%2!=0, F.col("xyz2")-F.col("xyz3")+1).otherwise(F.col("xyz2")-F.col("xyz3")))\
.withColumn("xyz4", (F.col("xyz9")/2))\
.withColumn("xyz6", F.col("xyz4")+F.col("xyz3"))\
.withColumn("xyz7", F.when(F.col("xyz10")%2==0,(F.col("xyz4")+F.col("xyz3")+1)).otherwise(F.lit(None)))\
.withColumn("xyz5", F.row_number().over(w))\
.withColumn("medianr", F.when(F.col("xyz6")==F.col("xyz5"), F.col("xyz")).when(F.col("xyz7")==F.col("xyz5"),F.col("xyz")).otherwise(F.lit(None)))\
.withColumn("medianr2", (F.mean("medianr").over(w2)))\
.withColumn("xyz", F.when(F.col("xyz").isNull(), F.col("medianr2")).otherwise(F.col("xyz")))\
.select("I_id","p_id","xyz")\
.orderBy("I_id").show()

+----+----+---+
|I_id|p_id|xyz|
+----+----+---+
|   1|   5|  4|
|   1|   5|  1|
|   1|   5|  4|
|   1|   5|  4|
|   2|   5|  2|
|   2|   5|  2|
|   2|   5|  1|
|   2|   5|  2|
|   2|   5|  4|
+----+----+---+