Pyspark: How to remove an item from a collect_set? - apache-spark

Question

In the following dataframe:
from pyspark.sql import functions as F
df = sqlContext.createDataFrame([
("a", "code1", "name"),
("a", "code1", "name2"),
("a", "code2", "name2"),
], ["id", "code", "name"])
df.show()
You can run this command to get a list of distinct values:
df.groupby("id").agg(F.collect_set("code")).show()
+---+-----------------+
| id|collect_set(code)|
+---+-----------------+
| a| [code2, code1]|
+---+-----------------+
How do you remove an item in the above collect_set? E.g. how to remove 'code2'

Answer 1

Update for Spark 2.4+: You can achieve this with array_remove:
df_grouped = df.groupby("id")\
.agg(F.array_remove(F.collect_set("code"), "code2").alias("codes"))
Original answer for Spark 2.3 and below
AFAIK there is no way to dynamically iterate over an ArrayType(), so if your data is already in an array you have two options:
Option 1: Explode, Filter, Collect
Use pyspark.sql.functions.explode() to turn the elements of the array into separate rows. Then use pyspark.sql.DataFrame.where() to filter out the desired values. Finally do a groupBy() and collect_set() to gather the data back into one row.
df_grouped = df.groupby("id").agg(F.collect_set("code").alias("codes"))
df_grouped.select("*", F.explode("codes").alias("exploded"))\
.where(~F.col("exploded").isin(["code2"]))\
.groupBy("id")\
.agg(F.collect_set("exploded").alias("codes"))\
.show()
#+---+-------+
#| id| codes|
#+---+-------+
#| a|[code1]|
#+---+-------+
Option 2: Use a UDF
def filter_code(array):
bad_values={"code2"}
return [x for x in array if x not in bad_values]
filter_code_udf = F.udf(lambda x: filter_code(x), ArrayType(StringType()))
df_grouped = df.groupby("id").agg(F.collect_set("code").alias("codes"))
df_grouped.withColumn("codes_filtered", filter_code_udf("codes")).show()
#+---+--------------+--------------+
#| id| codes|codes_filtered|
#+---+--------------+--------------+
#| a|[code2, code1]| [code1]|
#+---+--------------+--------------+
Of course, if you are starting from your original dataframe (before the groupBy() and collect_set()) you can just filter the desired values first:
df.where(~F.col("code").isin(["code2"])).groupby("id").agg(F.collect_set("code")).show()
#+---+-----------------+
#| id|collect_set(code)|
#+---+-----------------+
#| a| [code1]|
#+---+-----------------+