如何在Pandas列中自然处理Enum类的值？

你想要利用IntEnum的名称-数值映射，在Pandas分类列中显示友好的标签，同时保留数值存储的效率，这确实是个实用的需求。你之前的尝试直接将Enum成员作为分类，导致打印时只显示数值，转字符串又会丢失数值比较的能力，这里提供一个精准的解决方案：

核心思路 #

让Pandas分类列底层存储Enum的数值，同时将分类标签映射为Enum的名称，这样既保证存储效率，又能显示自然标签，还支持用Enum成员直接做比较。

实现代码 #

from enum import IntEnum
import pandas as pd

class CategoryType(IntEnum):
    RED = 0
    YELLOW = 1
    GREEN = 2

# 原始数值列表
raw_values = [2, 2, 1, 2, 0]

# 创建分类列：指定分类为Enum的数值，确保底层存储效率
cat = pd.Categorical(
    raw_values,
    categories=[member.value for member in CategoryType],
    ordered=False  # 根据需求决定是否需要有序分类
)

# 将分类数值重命名为对应的Enum名称，实现显示标签
cat = cat.rename_categories({member.value: member.name for member in CategoryType})

# 转为DataFrame查看效果
cat_data = pd.DataFrame(cat, columns=["category"])
print(cat_data)

运行后输出的DataFrame会显示：

category
0     GREEN
1     GREEN
2    YELLOW
3     GREEN
4       RED

验证数值比较功能 #

你关心的cat_data == CategoryType.RED这类操作完全可以正常执行：

# 筛选出RED对应的行
print(cat_data[cat_data["category"] == CategoryType.RED])

输出结果：

category
4       RED

原理说明 #

• 底层存储：分类列的codes属性仍然是原始数值的编码，没有额外的字符串存储开销，保证了效率。
• 显示逻辑：通过rename_categories将数值映射为Enum名称，Pandas显示时会使用这些名称，但实际比较时仍基于底层数值，所以Enum成员（本质是数值类型）可以直接参与比较。

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