使用 window 函数将 NA 替换为 pyspark 中的中位数

replace NA with median in pyspark using window function

我想在 pyspark 中使用 window 函数根据分区列用 medain 替换 NA?

示例输入:

所需输出:

正在创建您的数据框:

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

为了使解决方案尽可能通用和动态,我必须创建许多新列来计算中位数,并能够将其发送到空值。话虽如此,解决方案不会很慢,并且可以针对大数据进行扩展。

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