Pyspark collect list

user_8832641 0
Viewed 2266

I am doing a group by over a column in a pyspark dataframe and doing a collect list on another column to get all the available values for column_1. As below.

Column_1 Column_2
A        Name1
A        Name2
A        Name3
B        Name1
B        Name2
C        Name1
D        Name1
D        Name1
D        Name1
D        Name1

The output that i get is a collect list of column_2 with column_1 grouped.

Column_1 Column_2
A        [Name1,Name2,Name3]  
B        [Name1,Name2]
C        [Name1]
D        [Name1,Name1,Name1,Name1]

Now when all the values within collect list are same, i just want to display it only once and not four times. Below is the expected output.

Expected output:

Column_1 Column_2
A        [Name1,Name2,Name3]  
B        [Name1,Name2]
C        [Name1]
D        [Name1]

Is there a way to do this in pyspark?

pyspark apache-spark dataframe
2 Answers

Use collect_set to eliminate duplicates and use array_sort(From Spark-2.4.0) to sort the arrays.

  • (or) use array_distinct(From Spark-2.4.0) to eliminate duplicates from collect_list.
df.show()
#+--------+--------+
#|Column_1|Column_2|
#+--------+--------+
#|       A|   Name1|
#|       A|   Name2|
#|       A|   Name3|
#|       B|   Name1|
#|       B|   Name2|
#|       C|   Name1|
#|       D|   Name1|
#|       D|   Name1|
#+--------+--------+

from pyspark.sql.functions import *    
df.groupBy('Column_1').agg(array_sort(collect_set(col('Column_2'))).alias("Column_2")).orderBy("Column_1").show(10,False)

#using array_distinct,array_sort functions   
df.groupBy('Column_1').agg(array_sort(array_distinct(collect_list(col('Column_2')))).alias("Column_2")).orderBy("Column_1").show(10,False)
#+--------+---------------------+
#|Column_1|Column_2             |
#+--------+---------------------+
#|A       |[Name1, Name2, Name3]|
#|B       |[Name1, Name2]       |
#|C       |[Name1]              |
#|D       |[Name1]              |
#+--------+---------------------+
user_76326950

Adding to my comments above, if order matters in order of occurance (check input df):

+--------+--------+
|Column_1|Column_2|
+--------+--------+
|A       |Name1   |
|A       |Name3   | <-Name3 occurs first
|A       |Name2   |
|B       |Name1   |
|B       |Name2   |
|C       |Name1   |
|D       |Name1   |
|D       |Name1   |
|D       |Name1   |
|D       |Name1   |
+--------+--------+

you can first assign an index and the drop duplicates + collect list:

(df.withColumn("idx",F.monotonically_increasing_id()).dropDuplicates(["Column_1","Column_2"])
.orderBy("idx").groupby("Column_1").agg(F.collect_list("Column_2").alias("Column_2"))
 .orderBy("Column_1")).show(truncate=False)

+--------+---------------------+
|Column_1|Column_2             |
+--------+---------------------+
|A       |[Name1, Name3, Name2]|
|B       |[Name1, Name2]       |
|C       |[Name1]              |
|D       |[Name1]              |
+--------+---------------------+
user_98406370
Related