Conditionally fill NA based on aggregation of corresponding column values

user_8145971 0
Viewed 51

I have a dataframe with 3 columns- brand, color, and rating:

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

Expected output:

The average rating for all non-null Lacoste items is (6+5+3+9+4)/5 = 5.4; all null Rating values with the Lacoste brand should be set to 5.4.

The average rating for all non-null Gap items is (8+3+5+5)/4 = 5.25; all null Rating values with the Gap brand should be set to 5.25.

There is no average rating for non-null Banana Republic items, so we instead will aggregate by Color; the average rating for all non-null orange items is (8+5+9+5)/4 = 6.75, so we will set the null Banana Republic item, which is orange, to 6.75.

This is my attempt to use a conditional statement to fill in the value from one computed dataframe for a null in this dataframe:

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

I'm struggling with even a very basic first step of isolating a single record to conditionally fill it.

pyspark apache-spark-sql apache-spark python
1 Answers

You can use mean over a partition of brand/color, and use coalesce to replace nulls with the mean until it finds a mean that is not null.

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