PySpark分组聚合Type列：处理空值与修正去重结果问题

先处理Type列的空元素问题 #

首先得把Type列里的NULL、空串，还有拆分后产生的空值都清掉，不然聚合后会残留空元素：

• 用when+trim判断，把NULL或空串直接转成空数组
• 对非空的Type值，用split按逗号拆分后，再用array_remove删掉拆分出来的空元素

代码片段：

from pyspark.sql import functions as F

# 清洗Type列，移除所有空相关的无效值
cleaned_df = df.withColumn(
    "clean_type",
    F.when(
        F.trim(F.col("Type")).isNull() | (F.trim(F.col("Type")) == ""),
        F.array()
    ).otherwise(
        F.array_remove(F.split(F.col("Type"), ","), "")
    )
)

解决Data3为NULL时单独成组的问题 #

PySpark里分组键是NULL的话会各自单独成组，要让所有Data3为NULL的行归为同一组，得先把NULL替换成一个统一的标识（比如__NULL__），分组完再替换回NULL：

# 给NULL的Data3设置统一分组标识
group_ready_df = cleaned_df.withColumn(
    "group_data3",
    F.coalesce(F.col("Data3"), F.lit("__NULL__"))
)

# 分组聚合，合并并去重Type元素
result_df = group_ready_df.groupBy(
    "Data1", "Data2", "group_data3"
).agg(
    F.array_distinct(F.flatten(F.collect_list("clean_type"))).alias("unique_types")
).withColumn(
    "Data3",
    F.when(F.col("group_data3") == "__NULL__", F.lit(None)).otherwise(F.col("group_data3"))
).drop("group_data3")

完整可运行示例 #

from pyspark.sql import SparkSession
from pyspark.sql import functions as F

# 模拟测试数据
spark = SparkSession.builder.appName("TypeMergeFix").getOrCreate()
test_data = [
    ("2024-01-01", "A,B", 1, "X", None),
    ("2024-01-02", ",C,,", 1, "X", None),
    ("2024-01-03", None, 1, "X", 5),
    ("2024-01-04", "B,D", 1, "X", None),
    ("2024-01-05", "", 1, "X", 5)
]
df = spark.createDataFrame(test_data, ["Date", "Type", "Data1", "Data2", "Data3"])

# 清洗Type列
cleaned_df = df.withColumn(
    "clean_type",
    F.when(
        F.trim(F.col("Type")).isNull() | (F.trim(F.col("Type")) == ""),
        F.array()
    ).otherwise(
        F.array_remove(F.split(F.col("Type"), ","), "")
    )
)

# 处理分组键NULL并完成聚合
result_df = cleaned_df.withColumn(
    "group_data3",
    F.coalesce(F.col("Data3"), F.lit("__NULL__"))
).groupBy(
    "Data1", "Data2", "group_data3"
).agg(
    F.array_distinct(F.flatten(F.collect_list("clean_type"))).alias("unique_types")
).withColumn(
    "Data3",
    F.when(F.col("group_data3") == "__NULL__", F.lit(None)).otherwise(F.col("group_data3"))
).drop("group_data3")

# 查看结果
result_df.show(truncate=False)

核心逻辑说明 #

• array_remove：彻底清理拆分后Type数组里的空元素，避免聚合后出现无效空值
• coalesce替换NULL：让所有Data3为NULL的行被分到同一组，解决单独成组的问题
• flatten+array_distinct：先把收集到的多个Type数组合并成一个大数组，再去重，得到最终的唯一Type列表

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