PySpark使用含pandas.NA键的字典替换列值报错如何解决

错误根因 #

报错是因为PySpark的lit()函数无法直接序列化pandas.NA（NAType类型对象），你原代码中直接把包含pandas.NA的字典键传入create_map构造映射表达式，导致执行失败。

修改方案 #

我们把字典中的空键映射和普通键映射拆分处理，空值匹配逻辑单独实现，不放入create_map构造逻辑中，修改后的recode函数如下：

import pandas as pd
from pyspark.sql import Column
from pyspark.sql.functions import col, create_map, lit, when, isnull
from itertools import chain

def recode(col_name, map_dict, default=None):
    if not isinstance(col_name, Column): # 支持传入列名字符串或列实例
        col_name = col(col_name)
    
    # 拆分空键映射和普通键映射
    na_mapping = None
    normal_map = {}
    for k, v in map_dict.items():
        if pd.isna(k):
            na_mapping = v
        else:
            normal_map[k] = v
    
    # 仅用普通键构造映射表达式
    mapping_expr = create_map([lit(x) for x in chain(*normal_map.items())])
    map_result = mapping_expr.getItem(col_name)
    
    # 拼接匹配逻辑：优先匹配空值，再匹配普通映射，最后走默认值
    result = when(isnull(col_name), na_mapping)
    if default is None:
        result = result.when(~isnull(map_result), map_result)
    else:
        result = result.when(~isnull(map_result), map_result).otherwise(default)
    return result

调用逻辑不需要修改，直接使用原有代码即可得到预期输出：

data = df.withColumn("column", recode('column', swap, default=0))

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