如何将DataFrame中Enum列的单元格值以Enum类型获取？

你遇到的问题是Polars的pl.Enum列默认返回字符串而非Enum实例，这是因为Polars的Enum本质是分类类型，底层用整数编码存储，但对外默认输出字符串形式。以下是几种获取Enum类型值的方法：

1. 先定义对应Python Enum类 #

首先需要创建和Polars Enum分类匹配的Python标准库Enum类：

import polars as pl
from enum import Enum

# 定义Python Enum类
class Flags(Enum):
    FOO = "foo"
    BAR = "bar"
    BAZ = "baz"

# 初始化Polars DataFrame
pl_flags = pl.Enum(["foo", "bar", "baz"])
df = pl.DataFrame(data={"x": ["foo", "bar", "baz"]}, schema_overrides={"x": pl_flags})

2. 方法一：通过底层编码映射Enum #

利用Polars Enum的物理存储值（整数索引）转换为Enum实例：

# 获取列的底层整数编码值
physical_values = df["x"].to_physical().to_list()
# 获取Polars Enum的分类列表
categories = df["x"].dtype.categories
# 映射为Python Enum实例
enum_values = [Flags(categories[i]) for i in physical_values]

# 验证类型
print([type(v) for v in enum_values])  # 输出 [<enum 'Flags'>, <enum 'Flags'>, <enum 'Flags'>]

3. 方法二：直接从字符串转换Enum #

如果不需要底层编码，直接将字符串值映射到Enum：

string_values = df["x"].to_list()
enum_values = [Flags(v) for v in string_values]

print([type(v) for v in enum_values])  # 输出 [<enum 'Flags'>, <enum 'Flags'>, <enum 'Flags'>]

4. 方法三：在DataFrame中批量转换为Enum列 #

使用Polars的map_elements函数，直接在DataFrame中生成存储Enum实例的列：

df = df.with_columns(
    pl.col("x").map_elements(lambda s: Flags(s), return_dtype=pl.Object).alias("x_enum")
)

# 验证单个值类型
print(type(df[0, "x_enum"]))  # 输出 <enum 'Flags'>

补充说明 #

Polars的pl.Enum是为高效存储和处理分类数据设计的，默认返回字符串是为了简化日常操作。若需要Python Enum实例，必须手动通过上述方式完成映射，因为Polars不会自动将分类值转换为标准库的Enum类型。

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