如何在PySpark中使用字典替换数组列中的指定元素？

问题描述 #

我有一个名为fruits的列，每行数据为数组类型，格式示例：["apple","banana","mango"]。现有一个{旧值:新值}格式的字典，需要将数组中的apple替换为grapes，请问如何在PySpark中实现？

我的尝试 #

import pyspark.sql.functions as F

df_silver_3 = df_silver_2.filter(F.col("fruits").isNotNull())\
    .withColumn("fruits_cleaned", F.regexp_replace(df_silver_2["fruits"].cast("string"), r"[\[\]]", ""))\
    .select("fruits", "fruits_cleaned")

替换字典 #

replace_dict = {
    "apple": "grapes",  
    "value2": "replacement2"
}

尝试的替换代码 #

from pyspark.sql.functions import col, when

df_replaced = df_silver_3.select(
    [when(col("fruits_cleaned") == key, value).otherwise(col("fruits_cleaned")).alias("fruits_cleaned") for key, value in replace_dict.items()] + 
    [col("fruits")]  # 需要的话保留原列
)

当前尝试的问题 #

上述方案将数组转为字符串处理，再通过等值判断替换，无法实现数组内单个元素的替换——fruits_cleaned是逗号分隔的字符串（如"apple,banana,mango"），等值判断仅会匹配整个字符串等于apple的情况，完全无法处理数组内的元素替换逻辑。

正确实现方法 #

无需转换数组为字符串，直接针对数组类型列操作，使用transform遍历数组元素，结合映射表完成替换：

from pyspark.sql import functions as F

# 将Python字典转为Spark可识别的键值映射
replace_map = F.create_map([F.lit(x) for pair in replace_dict.items() for x in pair])

# 对fruits数组列执行元素替换
df_final = df_silver_2.filter(F.col("fruits").isNotNull())\
    .withColumn("fruits_replaced", F.transform("fruits", lambda x: F.coalesce(replace_map[x], x)))\
    .select("fruits", "fruits_replaced")

代码说明 #

1. create_map：将Python字典转换为Spark支持的键值映射结构，用于快速查找替换值；
2. transform：遍历fruits数组的每个元素，对元素执行替换逻辑——若元素在映射表中有对应新值则替换，否则保留原元素（coalesce处理元素不在映射表中的场景）。

预期结果 #

原数组["apple","banana","mango"]会被替换为["grapes","banana","mango"]

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