Spark DataFrame GroupBy collect_list and fill missing valeus with 0

user_2726038 0
Viewed 95

I have a dataframe of this kind

+------+--------+--------+
|    id|category|quantity|
+------+--------+--------+
|merch1|   fruit|    20.0|
|merch1| veggies|   300.0|
|merch1|   diary|    10.0|
|merch1|organics|    12.0|
|merch1|  frozen|    11.0|
|merch2|   fruit|     6.0|
|merch2|   diary|     6.0|
|merch2|  frozen|     8.0|
|merch3| veggies|    13.0|
|merch3|organics|     4.0|
|merch3|  frozen|    10.0|
|merch4|   fruit|    28.0|
|merch4|organics|    11.0|
+------+--------+--------+

I would like to groupBy id and generate a vector on quantity as an ordered collect_list, such that if the category is absent it fills with 0.0. The distinct categories are

+--------+                                                                      
|category|
+--------+
|  frozen|
|   diary|
| veggies|
|organics|
|   fruit|
+--------+

For example for merch1 since everything is present the vector for merch1 it is [11.0, 10.0, 300.0, 12.0, 20.0] for merch2 it is [8.0, 6.0, 0.0, 0.0, 6.0] for merch3 it is [10.0, 0.0, 13.0, 4.0, 0.0] for merch4 it is [0.0, 0.0, 0.0, 28.0, 11.0]

Hence the final dataframe I am looking for is

+------+--------+----------------------+
|    id| vector                        |
+------+--------+----------------------+
|merch1|[11.0, 10.0, 300.0, 12.0, 20.0]|
|merch2|[8.0, 6.0, 0.0, 0.0, 6.0]      |
|merch3|[10.0, 0.0, 13.0, 4.0, 0.0]    |
|merch4|[0.0, 0.0, 0.0, 28.0, 11.0]    |
+------+-------------------------------+
pyspark apache-spark
1 Answers

We can achieve this in two steps, convert rows to columns at group level(here id) using pivot and use array SQL function to create the list as expected.

from pyspark.sql import function as f
df.show()
+------+--------+--------+
|    id|category|quantity|
+------+--------+--------+
|merch1|   fruit|    20.0|
|merch1| veggies|   300.0|
|merch1|   diary|    10.0|
|merch1|organics|    12.0|
|merch1|  frozen|    11.0|
|merch2|   fruit|     6.0|
|merch2|   diary|     6.0|
|merch2|  frozen|     8.0|
|merch3| veggies|    13.0|
|merch3|organics|     4.0|
|merch3|  frozen|    10.0|
|merch4|   fruit|    28.0|
|merch4|organics|    11.0|
+------+--------+--------+  
    
df1 = df.groupby('id').pivot('category').agg(f.first('quantity')).fillna(0)
df1.show()
+------+-----+------+-----+--------+-------+
|    id|diary|frozen|fruit|organics|veggies|
+------+-----+------+-----+--------+-------+
|merch2|  6.0|   8.0|  6.0|     0.0|    0.0|
|merch4|  0.0|   0.0| 28.0|    11.0|    0.0|
|merch1| 10.0|  11.0| 20.0|    12.0|  300.0|
|merch3|  0.0|  10.0|  0.0|     4.0|   13.0|
+------+-----+------+-----+--------+-------+
    
df1.select('id',f.array(df1.columns[1:]).name('vector')).show(truncate=False)
+------+-------------------------------+
|id    |vector                         |
+------+-------------------------------+
|merch2|[6.0, 8.0, 6.0, 0.0, 0.0]      |
|merch4|[0.0, 0.0, 28.0, 11.0, 0.0]    |
|merch1|[10.0, 11.0, 20.0, 12.0, 300.0]|
|merch3|[0.0, 10.0, 0.0, 4.0, 13.0]    |
+------+-------------------------------+
user_73583920
Related