PySpark：无需转换字典替换DataFrame列值的方法

问题背景 #

现有字典：

d = {"animal": ["cat", "dog", "turtle"], "fruit" : ["banana", "apple"]}

以及Spark DataFrame：

+-----------+
|some_column|
+-----------+
|     banana|
|        cat|
|      apple|
|      other|
|       null|
+-----------+

期望输出结果：

+-----------+
|some_column|
+-----------+
|      fruit|
|     animal|
|      fruit|
|      other|
|       null|
+-----------+

已知把字典转成{"apple": "fruit", "banana": "fruit", ...}的键值对形式后，可以用df.na.replace实现替换，但希望不改动原字典结构完成需求。

解决方案 #

下面两种方法都不需要转换原字典，直接实现映射：

方法1：用when函数构建条件链 #

通过循环遍历原字典，为每个类别对应的元素列表构建匹配条件，生成条件表达式后替换字段：

from pyspark.sql import functions as F

# 初始化表达式，默认保留原字段值
expr = F.col("some_column")

# 遍历字典，逐个添加匹配规则
for category, items in d.items():
    expr = F.when(F.col("some_column").isin(items), category).otherwise(expr)

# 应用表达式得到结果
result_df = df.withColumn("some_column", expr)
result_df.show()

这个方法会依次检查字段值是否属于某个类别的元素列表，匹配成功就替换为类别名，不匹配则保留原内容（包括null和other）。

方法2：自定义UDF #

如果需要更复杂的匹配逻辑，自定义UDF是更灵活的选择：

from pyspark.sql import functions as F
from pyspark.sql.types import StringType

def map_value(val):
    if val is None:
        return None
    # 遍历字典找匹配的类别
    for category, items in d.items():
        if val in items:
            return category
    # 无匹配时返回原值
    return val

# 注册UDF
map_udf = F.udf(map_value, StringType())

# 应用UDF完成替换
result_df = df.withColumn("some_column", map_udf(F.col("some_column")))
result_df.show()

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