在withColumn中使用带有PySpark的"python枚举“的最佳方法(”myColumn“，myEnum.Giraffe)

from pyspark.sql import SparkSession
from pyspark.sql.functions import udf
from enum import Enum

# 创建SparkSession
spark = SparkSession.builder.getOrCreate()

# 定义枚举类
class MyEnum(Enum):
    Giraffe = 1
    Elephant = 2
    Lion = 3

# 创建UDF将枚举值转换为列值
enum_to_column = udf(lambda x: x.name if x else None)

# 创建DataFrame
data = [("Alice", MyEnum.Giraffe), ("Bob", MyEnum.Elephant), ("Charlie", MyEnum.Lion)]
df = spark.createDataFrame(data, ["Name", "Animal"])

# 使用withColumn和UDF将枚举值转换为列值
df = df.withColumn("AnimalName", enum_to_column(df.Animal))

# 显示结果
df.show()