在PySpark UDF的withColumn中使用Enum类遇验证错误求助

我有一个包含多列的DataFrame（inputDF），其中productcode是其中一列。我希望使用从其他目录导入的Enum类product_code（导入语句：from common.codes.product_code import product_code）验证该列的值，若productcode列的值不在Enum类定义的范围内，则在新列enum_column中写入False。但运行代码时出现错误提示：Error: '{enum_type}' is not a valid Enum class.，且执行globals().get(enum_type)时返回None，不清楚原因。

原代码如下：

from common.codes.product_code import product_code

def use_enum_in_method(enum_type,value):
    enum_class = globals().get(enum_type)
    if enum_class is None or not issubclass(enum_class, Enum):
        print(f"Error: '{enum_type}' is not a valid Enum class.")
        return False
    else:
        values_list = [member.value for member in enum_class]
        if value not in values_list:
            return False
        else:
            return True

validate_enum = udf(lambda a,value: use_enum_in_method(a,value), BooleanType())
df = inputDF.withColumn("enum_column",validate_enum(lit("product_code".lower()),"productcode"))
df.show()

Enum类定义：

from enum import Enum

class product_code(Enum):
    CEP = "CEP"
    COF = "COF"
    HMO = "HMO"
    MC = "MC"
    MCF = "MCF"
    MCR = "MCR"
    MCS = "MCS"
    MD = "MD"
    MDE = "MDE"
    MEP = "MEP"
    MLI = "MLI"
    MMO = "MMO"
    MMP = "MMP"
    MOS = "MOS"
    MP = "MP"
    MPO = "MPO"
    MR = "MR"
    MRB = "MRB"
    MRF = "MRF"
    MRS = "MRS"
    POS = "POS"
    PPO = "PPO"
    SN1 = "SN1"
    SN2 = "SN2"
    SN3 = "SN3"
    SN4 = "SN4"
    SN5 = "SN5"
    SN6 = "SN6"
    M99 = "M99"
    MCL = "MCL"

1. Spark UDF执行上下文隔离：UDF运行在Worker节点的独立进程中，这些进程的全局命名空间（globals()）和提交作业的主进程不共享，因此在UDF内部调用globals().get(enum_type)无法找到主进程中导入的product_code类。
2. 字符串传递Enum类不可行：通过字符串"product_code"传递类名到UDF，Worker节点无法将该字符串关联到实际的Enum类对象，必然返回None。

方案一：UDF内部直接导入Enum类（最简方式） #

确保Worker节点的Python环境能访问common.codes.product_code模块，直接在UDF内导入并使用Enum类：

from pyspark.sql import functions as F
from pyspark.sql.types import BooleanType
from enum import Enum
from common.codes.product_code import product_code

# 提前生成Enum值集合，避免每次调用重复遍历Enum
PRODUCT_CODE_VALUES = {member.value for member in product_code}

def validate_product_code(value):
    # 处理空值情况，避免报错
    return value in PRODUCT_CODE_VALUES if value is not None else False

# 定义UDF，直接关联验证函数
validate_enum_udf = F.udf(validate_product_code, BooleanType())

# 调用UDF处理DataFrame
df = inputDF.withColumn("enum_column", validate_enum_udf(F.col("productcode")))
df.show()

方案二：使用广播变量优化（适合大Enum集合） #

如果Enum值较多，可将值集合通过广播变量传递，减少Worker节点的内存重复占用：

from pyspark.sql import functions as F
from pyspark.sql.types import BooleanType
from common.codes.product_code import product_code
from pyspark.sql import SparkSession

spark = SparkSession.builder.getOrCreate()

# 生成Enum值集合并广播到所有Worker节点
product_code_broadcast = spark.sparkContext.broadcast({member.value for member in product_code})

def validate_product_code(value):
    return value in product_code_broadcast.value if value is not None else False

validate_enum_udf = F.udf(validate_product_code, BooleanType())

df = inputDF.withColumn("enum_column", validate_enum_udf(F.col("productcode")))
df.show()

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