如何在PySpark DataFrame中生成含空值列名列表的audit列

需求说明 #

给定如下DataFrame：

+---+--------+------+-----+-------+
| id|    name|  city|state|country|
+---+--------+------+-----+-------+
|  1|abhishek|  pune| null|   null|
|  2| devansh|  pune|   mh|  india|
|  3| urvashi|bglore|   ka|   null|
+---+--------+------+-----+-------+

需要新增audit列，存储每行中存在空值的列名列表：

• 第一行：[state,country]
• 第二行：[]（空列表）
• 第三行：[country]

PySpark 实现方案 #

利用isNull()、when()、array()和array_filter()函数组合实现：

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, array, when, array_filter

# 初始化Spark会话
spark = SparkSession.builder.appName("null_column_audit").getOrCreate()

# 构建示例DataFrame
data = [
    (1, "abhishek", "pune", None, None),
    (2, "devansh", "pune", "mh", "india"),
    (3, "urvashi", "bglore", "ka", None)
]
df = spark.createDataFrame(data, ["id", "name", "city", "state", "country"])

# 生成audit列的核心逻辑
audit_expr = array_filter(
    # 遍历所有列，空值列返回列名，非空列返回null
    array(*[when(col(c).isNull(), c) for c in df.columns]),
    # 过滤掉数组中的null值，保留空值列名
    lambda x: x.isNotNull()
)

# 添加audit列并展示结果
df_with_audit = df.withColumn("audit", audit_expr)
df_with_audit.show(truncate=False)

执行结果：

+---+--------+------+-----+-------+----------------+
|id |name    |city  |state|country|audit           |
+---+--------+------+-----+-------+----------------+
|1  |abhishek|pune  |null |null   |[state, country]|
|2  |devansh |pune  |mh   |india  |[]              |
|3  |urvashi |bglore|ka   |null   |[country]       |
+---+--------+------+-----+-------+----------------+

Pandas 实现方案 #

如果使用Pandas，可以通过apply()方法遍历每行生成列名列表：

import pandas as pd

# 构建示例DataFrame
data = [
    (1, "abhishek", "pune", None, None),
    (2, "devansh", "pune", "mh", "india"),
    (3, "urvashi", "bglore", "ka", None)
]
df_pd = pd.DataFrame(data, columns=["id", "name", "city", "state", "country"])

# 生成audit列
df_pd["audit"] = df_pd.apply(
    lambda row: [col_name for col_name in df_pd.columns if pd.isna(row[col_name])],
    axis=1
)

print(df_pd)

执行结果：

id      name    city state country             audit
0   1  abhishek    pune  None    None  [state, country]
1   2   devansh    pune    mh   india                []
2   3   urvashi  bglore    ka    None          [country]

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