pandas数据框中的枚举，无法在枚举列上进行分组吗？

Question

我刚刚学会了枚举（enums），觉得它们可以用在我正在编写的代码里。但是当我运行这段代码时，出现了错误。我是在做不该做的事情，还是这是个bug呢？

当我尝试对一个包含枚举的列进行分组时，出现了这个错误： 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)

我原本期待的输出是： 品牌 价格 BMW 45000 VOLVO 72000

Answer 1

另一种解决方案是为你的枚举定义顺序：

from enum import Enum
from functools import total_ordering

import pandas as pd


@total_ordering
class CarBrand(Enum):
    VOLVO = "Volvo"
    BMW = "BMW"

    def __lt__(self, other):
        return self.value < other.value


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)

输出结果是：

                price
brand                
CarBrand.BMW    45000
CarBrand.VOLVO  72000

Answer 2

pd.DataFrame.groupby 默认会进行排序。如果你不想排序，可以使用 sort=False 这个选项。

sum_per_brand = df.groupby('brand', sort=False).sum('price')

另外，你也可以使用一种支持排序的数据类型，比如 CategoricalDtype。