Pyspark：如何从collect_set中删除项目？

Question

在以下数据框中：

from pyspark.sql import functions as F
df = sqlContext.createDataFrame([
    ("a", "code1", "name"),
    ("a", "code1", "name2"),
    ("a", "code2", "name2"),
], ["id", "code", "name"])

df.show()

您可以运行此命令来获取不同值的列表：

df.groupby("id").agg(F.collect_set("code")).show()

+---+-----------------+
| id|collect_set(code)|
+---+-----------------+
|  a|   [code2, code1]|
+---+-----------------+

如何删除上述collect_set中的项目？例如如何删除'code2'

Answer 1

Spark 2.4+ 的更新：您可以通过以下方式实现此目的array_remove：

df_grouped = df.groupby("id")\
    .agg(F.array_remove(F.collect_set("code"), "code2").alias("codes"))

Spark 2.3 及以下版本的原始答案

AFAIK 无法动态迭代ArrayType()，因此如果您的数据已经在数组中，您有两个选择：

选项 1：分解、过滤、收集

用于pyspark.sql.functions.explode()将数组的元素转换为单独的行。然后使用pyspark.sql.DataFrame.where()过滤掉所需的值。最后执行groupBy()和collect_set()将数据收集回一行。

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]|
#+---+-------+

选项 2：使用 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]|
#+---+--------------+--------------+

当然，如果您从原始数据帧（在groupBy()和之前collect_set()）开始，您可以先过滤所需的值：

df.where(~F.col("code").isin(["code2"])).groupby("id").agg(F.collect_set("code")).show()
#+---+-----------------+
#| id|collect_set(code)|
#+---+-----------------+
#|  a|          [code1]|
#+---+-----------------+