如何在PySpark中按分组高效替换DataFrame的空值？

你之前拆分DataFrame再填充的方式会产生不必要的Shuffle开销，大数据量下效率很低，而且链式where的写法逻辑错误——第一个where已经过滤掉非A组的数据，后续where找不到B组数据，最终会得到空结果。

下面是两种高效的实现方式，无需拆分数据集：

单列空值替换 #

假设要处理的目标列是value，分组列是group，直接用when/otherwise组合判断分组和空值条件：

from pyspark.sql import functions as F

df = df.withColumn(
    "value",
    F.when(
        (F.col("group") == "A") & F.col("value").isNull(),
        F.lit(-999)
    ).when(
        (F.col("group") == "B") & F.col("value").isNull(),
        F.lit(0)
    ).otherwise(F.col("value"))
)

多列空值替换 #

如果需要批量处理多列，可以用循环或列表推导式批量生成填充逻辑：

from pyspark.sql import functions as F

# 指定需要处理的列
cols_to_fill = ["col1", "col2", "col3"]

# 循环处理每一列
for col_name in cols_to_fill:
    df = df.withColumn(
        col_name,
        F.when(
            (F.col("group") == "A") & F.col(col_name).isNull(),
            F.lit(-999)
        ).when(
            (F.col("group") == "B") & F.col(col_name).isNull(),
            F.lit(0)
        ).otherwise(F.col(col_name))
    )

这种方式全程在同一个DataFrame上操作，避免了拆分合并带来的性能损耗，内置函数的执行效率远高于拆分后单独处理的逻辑。

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