Pandas DataFrame枚举列执行groupby报错TypeError，求解决方案

问题详情 #

刚学习枚举（Enum），想在代码中使用，但对枚举列执行groupby时出现如下错误：

TypeError: '<' not supported between instances of 'CarBrand' and 'CarBrand'

出错代码：

import pandas as pd
from enum import Enum

class CarBrand(Enum):
    VOLVO = 'Volvo'
    BMW = 'BMW'

data = {
    'brand': [CarBrand.VOLVO,
              CarBrand.VOLVO, 
              CarBrand.BMW],
    'price': [35000, 
              37000, 
              45000]
}

df = pd.DataFrame(data)
sum_per_brand = df.groupby('brand').sum('price')
print(sum_per_brand)

期望输出：

brand  price
BMW    45000
VOLVO  72000

原因分析 #

Pandas的groupby操作需要对分组键进行排序或比较，而默认的Enum实例仅支持身份比较（判断是否为同一个实例），不支持<这类大小比较运算符，因此触发错误。

解决方案 #

方案1：使用枚举的value作为分组值 #

构造数据时直接存储枚举对应的实际值（而非枚举实例），将列数据转为字符串，即可正常执行groupby：

import pandas as pd
from enum import Enum

class CarBrand(Enum):
    VOLVO = 'Volvo'
    BMW = 'BMW'

data = {
    'brand': [CarBrand.VOLVO.value,
              CarBrand.VOLVO.value, 
              CarBrand.BMW.value],
    'price': [35000, 
              37000, 
              45000]
}

df = pd.DataFrame(data)
sum_per_brand = df.groupby('brand').sum('price')
print(sum_per_brand)

方案2：让枚举类继承str和Enum #

修改枚举类，使其同时继承str和Enum，这样枚举实例会具备字符串的比较能力，无需修改其他代码即可正常分组：

import pandas as pd
from enum import Enum

class CarBrand(str, Enum):
    VOLVO = 'Volvo'
    BMW = 'BMW'

data = {
    'brand': [CarBrand.VOLVO,
              CarBrand.VOLVO, 
              CarBrand.BMW],
    'price': [35000, 
              37000, 
              45000]
}

df = pd.DataFrame(data)
sum_per_brand = df.groupby('brand').sum('price')
print(sum_per_brand)

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