Filter values while using pyspark collect list

user_8832641 0
Viewed 755

I have the below pyspark dataframe.

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

I am using group by to aggregate the data and collect set to collect the data and remove duplicate values and creating the below output.

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

But my expected output is whenever X occurs with other values like Name1, Name2 etc. It has to be removed from the output set (Row A ,B). But whenever its the only value for a row, it has to be retained like row E.

Expected Output:

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

Is there a way to achieve this in pyspark?

pyspark apache-spark dataframe
1 Answers

From Spark-2.4 you can use array_remove function if the array size >1.

Example:

df.show(10,False)
#+--------+------------------------+
#|Column_1|Column_2                |
#+--------+------------------------+
#|A       |[Name1, Name2, Name3, X]|
#|E       |[X]                     |
#+--------+------------------------+

from pyspark.sql.functions import *
df.withColumn("Column_2",when(size(col("Column_2")) > 1 , array_remove(col("column_2"),'X')).otherwise(col("column_2"))).show(10,False)

#using filter function
df.withColumn("Column_2",when(size(col("Column_2")) > 1 , expr("filter(column_2, x -> not(x <=> 'X'))")).otherwise(col("column_2"))).show(10,False)
#+--------+---------------------+
#|Column_1|Column_2             |
#+--------+---------------------+
#|A       |[Name1, Name2, Name3]|
#|E       |[X]                  |
#+--------+---------------------+
user_76326950
Related