PySpark技术问答：如何从collect_set中移除指定元素？

嘿，这个需求我之前处理过，给你两种靠谱的方法来实现从collect_set集合里移除指定元素，就用你提供的示例DataFrame来演示：

首先先回顾下你的原始数据和基础聚合代码：

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

方法1：使用Spark内置的array_remove函数（推荐，Spark 2.4+支持） #

Spark 2.4及以上版本提供了array_remove函数，可以直接从数组（包括collect_set返回的集合数组）中移除指定元素，用法非常简洁：

df.groupby("id")
  .agg(F.array_remove(F.collect_set("code"), "code1").alias("filtered_codes"))
  .show()

执行后输出结果：

+---+--------------+
| id|filtered_codes|
+---+--------------+
|  a|       [code2]|
+---+--------------+

这个方法是Spark原生实现，性能比自定义UDF好很多，能满足大部分常规场景的需求。

方法2：自定义UDF（适配Spark 2.4以下版本） #

如果你的Spark版本低于2.4，没有array_remove函数，可以自定义一个UDF来实现过滤逻辑：

from pyspark.sql.types import ArrayType, StringType

# 定义一个用于移除指定元素的函数
def remove_target_element(arr, target):
    return [item for item in arr if item != target]

# 把函数注册为UDF
remove_element_udf = F.udf(remove_target_element, ArrayType(StringType()))

# 聚合后应用UDF移除code1
df.groupby("id")
  .agg(remove_element_udf(F.collect_set("code"), F.lit("code1")).alias("filtered_codes"))
  .show()

执行后得到的结果和方法1完全一致。如果需要更复杂的过滤规则（比如模糊匹配移除），也可以修改UDF里的逻辑来实现。

内容的提问来源于stack exchange，提问作者Micah Pearce