首页 > 解决方案 > 根据相应列值的聚合有条件地填充 NA

问题描述

我有一个包含 3 列的数据框 - 品牌、颜色和评级:

values = [('Lacoste', 'Red', 6), ('Gap', 'Orange', 8), ('Lacoste', 'Green', 5),
         ('Gap', 'Red', 3), ('Gap', 'Orange', 5), ('Lacoste', 'Green', 3),
         ('Lacoste', 'Orange', 9), ('Lacoste', 'Red', 4), ('Gap', 'Green', None),
         ('Lacoste', 'Red', None), ('Gap', 'Orange', 5), ('Lacoste', 'Green', None),
         ('Banana Republic', 'Orange', None)]
ratings = spark.createDataFrame(values, ['Brand', 'Color', 'Rating'])
ratings.show()

#+---------------+------+------+
#|          Brand| Color|Rating|
#+---------------+------+------+
#|        Lacoste|   Red|     6|
#|            Gap|Orange|     8|
#|        Lacoste| Green|     5|
#|            Gap|   Red|     3|
#|            Gap|Orange|     5|
#|        Lacoste| Green|     3|
#|        Lacoste|Orange|     9|
#|        Lacoste|   Red|     4|
#|            Gap| Green|  null|
#|        Lacoste|   Red|  null|
#|            Gap|Orange|     5|
#|        Lacoste| Green|  null|
#|Banana Republic|Orange|  null|
#+---------------+------+------+

预期输出:

所有非空 Lacoste 商品的平均评分为 (6+5+3+9+4)/5 = 5.4;Lacoste 品牌的所有 null Rating 值都应设置为 5.4。

所有非空 Gap 项目的平均评分为 (8+3+5+5)/4 = 5.25;Gap 品牌的所有 null Rating 值都应设置为 5.25。

非空香蕉共和国物品没有平均评分,因此我们将按颜色汇总;所有非空橙色项目的平均评分为 (8+5+9+5)/4 = 6.75,因此我们将设置为橙色的空香蕉共和国项目为 6.75。

这是我尝试使用条件语句从一个计算的数据帧中为该数据帧中的空值填充值:

brand_agg=ratings.groupBy("Brand").agg(avg("Rating").alias('Mean'))
brand_agg.show()

#+---------------+----+
#|          Brand|Mean|
#+---------------+----+
#|            Gap|5.25|
#|        Lacoste| 5.4|
#|Banana Republic|null|
#+---------------+----+

# this fails miserably
testing_df = ratings.withColumn('Rating', 
    when((ratings.Rating.isNull()) & 
    (brand_agg.Brand == ratings.Brand), 
    brand_agg.Mean).otherwise(ratings.Rating)) 

即使是非常基本的第一步,我也在努力隔离单个记录以有条件地填充它。

标签: pythonapache-sparkpysparkapache-spark-sql

解决方案


您可以使用mean品牌/颜色的分区,并用coalesce替换空值,mean直到找到不为空的平均值。

from pyspark.sql import functions as F, Window

filled = ratings.withColumn(
    'Rating',
    F.coalesce(
        F.col('Rating'), 
        F.mean('Rating').over(Window.partitionBy('Brand')), 
        F.mean('Rating').over(Window.partitionBy('Color'))
    )
)

filled.show()
+---------------+------+------+
|          Brand| Color|Rating|
+---------------+------+------+
|            Gap|Orange|   8.0|
|            Gap|Orange|   5.0|
|            Gap|Orange|   5.0|
|        Lacoste|Orange|   9.0|
|Banana Republic|Orange|  6.75|
|            Gap| Green|  5.25|
|        Lacoste| Green|   5.0|
|        Lacoste| Green|   3.0|
|        Lacoste| Green|   5.4|
|            Gap|   Red|   3.0|
|        Lacoste|   Red|   6.0|
|        Lacoste|   Red|   4.0|
|        Lacoste|   Red|   5.4|
+---------------+------+------+

推荐阅读