Best way to use a "python enum" with PySpark inside withColumn("myColumn", myEnum.Giraffe)

Question

I'd like to do something like this:

myEnum = Enum("myEnum", "Giraffe Crab Rat Human")

df_new = df.withColumn("AnimalEnumValueColumn", myEnum.Giraffe)

(I know this is incorrect, but I'd like an alternate solution that has the same intended effect)

user_14165730 · Answer

You need to use .value to get the value from the enum object:

import pyspark.sql.functions as F

df_new = df.withColumn("AnimalEnumValueColumn", F.lit(myEnum.Giraffe.value))

See the enum docs for more details.