PySpark中从多指标列提取去重值集合的实现问题

问题场景 #

现有一张包含ID、DATE及四个字符串指标列的表，示例输入：

需要新增total_ind列，该列为上述四个指标列的去重值集合，示例输出：

使用Python可实现，但用PySpark UDF（输入为拼接后的四列值）时出现内存不足错误，需寻求替代方案。

解决方案 #

避免使用Python UDF（序列化开销大，易引发内存问题），直接用PySpark内置函数实现，执行效率更高：

方法1：基础版 - array_distinct + array函数 #

直接将四列打包成数组后去重：

from pyspark.sql import functions as F

df = df.withColumn(
    "total_ind",
    F.array_distinct(
        F.array(
            F.col("primary_ind"),
            F.col("secondary_ind"),
            F.col("third_ind"),
            F.col("fourth_ind")
        )
    )
)

• 逻辑：array()把指定列转为数组，array_distinct()对数组元素去重，直接得到目标集合。

方法2：含空值处理版 #

如果指标列可能存在null，可先过滤空值再去重：

from pyspark.sql import functions as F

df = df.withColumn(
    "total_ind",
    F.array_distinct(
        F.filter(
            F.array("primary_ind", "secondary_ind", "third_ind", "fourth_ind"),
            lambda x: x.isNotNull()
        )
    )
)

• 逻辑：先用filter剔除数组中的空元素，再执行去重，避免集合包含无效空值。

方法3：可扩展版 - 先 unpivot 再聚合 #

如果后续需要新增更多指标列，可先将宽表转长表再聚合去重：

from pyspark.sql import functions as F

# 1. 把四列转为键值对格式（unpivot）
unpivot_df = df.select(
    "ID", "DATE",
    F.explode(
        F.array(
            F.struct(F.lit("primary").alias("type"), F.col("primary_ind").alias("code")),
            F.struct(F.lit("secondary").alias("type"), F.col("secondary_ind").alias("code")),
            F.struct(F.lit("third").alias("type"), F.col("third_ind").alias("code")),
            F.struct(F.lit("fourth").alias("type"), F.col("fourth_ind").alias("code"))
        )
    ).alias("ind")
).select("ID", "DATE", "ind.code")

# 2. 按ID、DATE聚合，收集去重后的code集合
result_df = unpivot_df.groupBy("ID", "DATE").agg(
    F.collect_set("code").alias("total_ind")
)

# 3. 关联回原表（保留原指标列）
final_df = df.join(result_df, on=["ID", "DATE"], how="left")

• 优势：指标列数量增加时只需修改unpivot部分，扩展性更强；适合需要对指标列做额外处理的场景。

为什么UDF会内存不足？ #

Python UDF需要在JVM和Python进程间反复做数据序列化/反序列化，处理大量数据时会产生额外内存开销；而PySpark内置函数基于JVM实现，执行效率高、内存占用低，能从根源避免这类问题。

内容的提问来源于stack exchange，提问作者Theo Mak