根据对应列值的聚合有条件地填充 NA

时间:2021-01-25 20:57:09

标签: python apache-spark pyspark apache-spark-sql pyspark-dataframes

我有一个包含 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 品牌的所有空评级值都应设置为 5.4。

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

非空 Banana Republic 商品没有平均评分,因此我们将按颜色汇总;所有非空橙色物品的平均评分为 (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)) 

即使是隔离单个记录以有条件地填充它的非常基本的第一步,我都在挣扎 - 感谢任何指导!

1 个答案:

答案 0 :(得分:1)

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

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