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'。

- Micah Pearce

---

可能是从PySpark DataFrame中的Python列表中删除元素的重复问题。 - pault

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及以下版本的原始答案

据我所知，没有办法动态地迭代一个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]|
#+---+-------+

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

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

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

- pault

---

有没有“isnotin”命令，我想从许多中过滤掉一个。无论如何，另外两个选项似乎都可以工作。谢谢！ - Micah Pearce

@MicahPearce，你只需要执行isin()的反操作 - 参见此帖子上的答案。 - pault

啊，他就是这么做的。这很有道理。 - Micah Pearce

@MicahPearce 还值得注意的是，如果可能的话，使用 udf 通常会更慢，应该避免使用。 - pault