按分組平均值填充 PySpark 數據框列的空值

Question

我有一個包含這兩列的大型 PySpark 數據框：

高速公路	速度_公里
路	70
服務	30
路	空值
路	70
服務	空值

我想通過該highway類別的平均值填充null值。

我試過用 groupBy 創建另一個數據框，最后得到了這個第二個：

高速公路	平均（速度_公里）
路	65
服務	30

但是我不知道如何使用它來僅填充null值，並且如果它們存在則不會丟失原始值。

第一個表的預期結果是：

高速公路	速度_公里
路	70
服務	30
路	65
路	70
服務	30

Answer 1

結合case when和由highway划分的窗函數可以很容易地解決這個問題。

from pyspark.sql import functions as F
from pyspark.sql import Window as W

(df
    .withColumn('speed_kph', F
        .when(F.col('speed_kph').isNull(), F.mean('speed_kph').over(W.partitionBy('highway')))
        .otherwise(F.col('speed_kph'))
    )
    .show()
)

# Output
# +-------+---------+
# |highway|speed_kph|
# +-------+---------+
# |   Road|     70.0|
# |   Road|     70.0|
# |   Road|     70.0|
# |Service|     30.0|
# |Service|     30.0|
# +-------+---------+