PySpark中使用collect_list时如何保留null值？

嘿，这个问题我太熟了！PySpark的collect_list和collect_set默认会悄悄过滤掉null值，确实挺让人头疼的，尤其是当你需要完整保留原始数据里的空值时。不过别担心，有几种靠谱的办法能解决这个问题，我给你一一拆解：

这个方法利用了array函数会保留null的特性，先把目标列包装成单元素数组，再用collect_list收集这些数组，最后用flatten把嵌套数组展开成普通列表——完美保留null值，而且是内置函数，性能比自定义UDF好太多！

from pyspark.sql.functions import array, collect_list, flatten

# 假设你要按cId和eId分组，收集amount和city列的列表并保留null
df_result = df.groupBy("cId", "eId") \
    .agg(
        flatten(collect_list(array("amount"))).alias("amount_list"),
        flatten(collect_list(array("city"))).alias("city_list")
    )

如果你的Spark版本低于2.4，没有flatten函数，那就写个简单的自定义UDF来实现。核心思路就是直接把分组后的迭代器转换成列表，不会过滤null：

from pyspark.sql.functions import udf
from pyspark.sql.types import ArrayType, StringType, DoubleType

# 针对不同数据类型写对应的UDF（比如amount是Double，city是String）
def collect_with_nulls_double(iterable):
    return list(iterable)

def collect_with_nulls_str(iterable):
    return list(iterable)

# 注册UDF，指定返回的数组类型
collect_nulls_double_udf = udf(collect_with_nulls_double, ArrayType(DoubleType()))
collect_nulls_str_udf = udf(collect_with_nulls_str, ArrayType(StringType()))

# 分组聚合
df_result = df.groupBy("cId", "eId") \
    .agg(
        collect_nulls_double_udf(collect_list("amount")).alias("amount_list"),
        collect_nulls_str_udf(collect_list("city")).alias("city_list")
    )

注意：UDF需要和目标列的数据类型匹配，不然会抛出类型错误哦。

如果以上两种方法都用不了，还可以用个小技巧：先把null替换成一个唯一的特殊标记，收集后再把标记转回去。但这个方法有局限性——如果你的列本身就包含这个标记，就会出错，所以只建议应急使用：

from pyspark.sql.functions import collect_list, concat_ws, split, when, col

# 选一个数据里不会出现的分隔符和null标记
SEP = "|||"
NULL_MARKER = "__SPECIAL_NULL__"

df_result = df.groupBy("cId", "eId") \
    .agg(
        split(
            concat_ws(SEP, collect_list(when(col("amount").isNull(), NULL_MARKER).otherwise(col("amount")))),
            SEP
        ).alias("amount_list"),
        split(
            concat_ws(SEP, collect_list(when(col("city").isNull(), NULL_MARKER).otherwise(col("city")))),
            SEP
        ).alias("city_list")
    ) \
    # 把标记转回到null
    .withColumn("amount_list", when(col("amount_list") == NULL_MARKER, None).otherwise(col("amount_list"))) \
    .withColumn("city_list", when(col("city_list") == NULL_MARKER, None).otherwise(col("city_list")))

最后再啰嗦一句：优先用方法1，内置函数的性能和稳定性都比UDF强，代码也更简洁。

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