Pyspark - Groupby and collect list over multiple columns and create multiple columns

user_8832641 0
Viewed 3043

I have the below dataframe over which I am trying to group by and aggregate data. 

Column_1 Column_2 Column_3
A        N1       P1
A        N2       P2
A        N3       P3
B        N1       P1
C        N1       P1
C        N2       P2

Required Output:

Column_1 Column_2 Column_3
A        N1,N2,N3 P1,P2,P3
B        N1       P1
C        N1,N2    P1,P2

I am able to do it over one column by creating a window using partition and groupby. Then I use collect list and group by over the window and aggregate to get a column. THis works for one column.

How to perform the same over 2 columns. Kindly help

pyspark apache-spark dataframe python
1 Answers

The agg function of the group by can take more than one aggreation function. You can add collect_list twice:

df.groupby('Column_1').agg(F.collect_list('Column_2'), F.collect_list('Column_3')).orderBy('Column_1').show()

prints

+--------+----------------------+----------------------+
|Column_1|collect_list(Column_2)|collect_list(Column_3)|
+--------+----------------------+----------------------+
|       A|          [N1, N2, N3]|          [P1, P2, P3]|
|       B|                  [N1]|                  [P1]|
|       C|              [N1, N2]|              [P1, P2]|
+--------+----------------------+----------------------+

For a simple grouping there is no need to use a Window.

user_21298010
Related