PySpark：如何用字典替换列值，避免重复调用when函数？

问题原因 #

你当前的循环逻辑完全错误：

• 每次循环都基于原始custoDF重新生成custoDF1，之前循环的修改完全被覆盖
• 每个when的otherwise直接返回当前循环的month，最后一次循环是num=12、month=Dec，所以不管Nummes是什么值，要么匹配12返回Dec，否则也返回Dec，最终结果全是Dec

正确实现方式 #

Spark完全支持用字典实现批量列值替换，不需要重复写12次when，推荐两种简洁写法：

方法1：链式拼接when-otherwise #

通过遍历字典，逐个拼接when条件，最后统一处理默认值：

from pyspark.sql.functions import col, when

months = {'1': 'Jan', '2': 'Feb', '3': 'Mar', '4': 'Apr', '5': 'May', '6': 'Jun',
         '7': 'Jul', '8': 'Aug', '9': 'Sep', '10':'Oct', '11': 'Nov', '12':'Dec'}

# 初始化第一个条件
month_expr = when(col("Nummes") == list(months.keys())[0], list(months.values())[0])
# 遍历剩余键值对，拼接when
for num, month in list(months.items())[1:]:
    month_expr = month_expr.when(col("Nummes") == num, month)
# 处理匹配不到的情况，这里设为None，可根据需求修改
month_expr = month_expr.otherwise(None)

custoDF1 = custoDF.withColumn("Month", month_expr)
custoDF1.select(col('Nummes').alias('NumMonth'), 'Month').distinct().orderBy("NumMonth").show(200)

方法2：用create_map实现字典映射（更高效） #

利用Spark的create_map函数直接构建映射表，通过键查找值，代码更简洁：

from pyspark.sql.functions import col, create_map, lit
from itertools import chain

months = {'1': 'Jan', '2': 'Feb', '3': 'Mar', '4': 'Apr', '5': 'May', '6': 'Jun',
         '7': 'Jul', '8': 'Aug', '9': 'Sep', '10':'Oct', '11': 'Nov', '12':'Dec'}

# 将字典转为(键, 值)的lit对，传入create_map
month_map = create_map(list(chain(*[(lit(num), lit(month)) for num, month in months.items()])))

custoDF1 = custoDF.withColumn("Month", month_map[col("Nummes")])
custoDF1.select(col('Nummes').alias('NumMonth'), 'Month').distinct().orderBy("NumMonth").show(200)

说明 #

这两种方法都能避免重复调用when，你的问题是之前的循环逻辑错误，并非Spark函数限制。

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