PySpark字典映射替换列值函数失效问题排查求助

先拆解你的代码里的核心问题，逐一说明：

关键错误点 #

• 变量名不匹配：函数参数是categorical_columns，但循环里用了未定义的cat_cols_list；参数是mapping，但判断条件里用了mapping_table，直接导致映射逻辑根本没触发。
• UDF闭包陷阱：循环内定义的lambda引用colmn，但Spark UDF是延迟执行的，等到实际计算时colmn已经是循环的最后一个值，所有列都会误用最后一列的映射规则。
• 映射表达式写法错误：注释里的mapping[colmn][colmn]不符合Spark列表达式语法，无法正确取到映射值。
• 未更新原DataFrame列表：你修改了循环里的dataframe变量，但最后没把它存回dataframes[i]，导致所有修改都没保留（注释里的赋值被你注释掉了）。

修复后的代码 #

下面是修正后的版本，同时改用Spark内置函数替代UDF（性能更好，避免闭包问题）：

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

def apply_hierarchy_masking(dataframes, mapping, categorical_columns, apply_mask_flag=0):
    if apply_mask_flag == 1:
        for i, dataframe in enumerate(dataframes):
            # 遍历传入的分类列参数，而非未定义的cat_cols_list
            for colmn in categorical_columns:
                # 检查当前列是否在映射表的key中
                if colmn in mapping:
                    # 先转成字符串类型（和原逻辑保持一致）
                    dataframe = dataframe.withColumn(colmn, dataframe[colmn].cast(T.StringType()))
                    
                    # 用Spark内置create_map构建映射表达式，替代UDF
                    mapping_pairs = []
                    for key, value in mapping[colmn].items():
                        mapping_pairs.append(F.lit(key))
                        mapping_pairs.append(F.lit(value))
                    mapping_expr = F.create_map(*mapping_pairs)
                    
                    # 用映射表达式替换原列，找不到匹配值则保留原值
                    dataframe = dataframe.withColumn(
                        colmn,
                        F.coalesce(mapping_expr[dataframe[colmn]], dataframe[colmn])
                    )
                else:
                    print(f"Not applying masking for column: {colmn}")
            
            # 必须把修改后的DataFrame存回原列表，否则修改无效
            dataframes[i] = dataframe
    else:
        print("Not applying hierarchy masking to the given data")
    return dataframes

修复说明 #

1. 统一变量名：使用函数传入的categorical_columns和mapping参数，避免未定义变量导致的逻辑失效。
2. 替换UDF为内置函数：用create_map构建映射规则，彻底避开闭包陷阱，同时Spark内置函数的执行效率远高于自定义UDF。
3. 保留未匹配值：用coalesce保证原列值不在映射字典中时，会保留原始值，和你原本的逻辑一致。
4. 更新原列表：将修改后的DataFrame赋值回dataframes[i]，确保修改能被保留并返回。

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