Pandas DataFrame - 对 dtype 为 'category' 的列进行聚合会导致性能缓慢

Question

我处理大型数据框，内存使用率很高。我读到如果我更改重复值列的dtype，可以节省大量内存。我试过了，确实减少了25%的内存使用率，但是我遇到了性能下降的问题，无法理解原因。我对'dtype'为'category'的列进行分组聚合，在更改'dtype'之前，它需要约1秒钟，而在更改后，需要约1分钟。以下代码演示了性能下降2倍的情况：

import pandas as pd
import random

animals = ['Dog', 'Cat']
days = ['Sunday', 'Monday', 'Tuesday', 'Wednesday', 'Thursday', 'Friday','Saturday']

columns_dict = {'animals': [],
                'days': []}

for i in range(1000000):
    columns_dict['animals'].append(animals[random.randint(0, len(animals)-1)])
    columns_dict['days'].append(days[random.randint(0, len(days)-1)])

# df without 'category' dtype
df = pd.DataFrame(columns_dict)

df.info(memory_usage='deep') # will result in memory usage of 95.5 MB

%timeit -n100 df.groupby('days').agg({'animals': 'first'})
# will result in: 100 loops, best of 3: 54.2 ms per loop

# df with 'category' dtype
df2 = df.copy()
df2['animals'] = df2['animals'].astype('category')

df2.info(memory_usage='deep') # will result in memory usage of 50.7 MB

%timeit -n100 df2.groupby('days').agg({'animals': 'first'})
# will result in: 100 loops, best of 3: 111 ms per loop

我试图了解的是这种缓慢的原因，以及是否有一种方法可以克服它。

谢谢！

- Lior Zamir

---

我在分类列和group_by操作方面遇到了完全相同的问题。请参见我的最近一篇SO问题。将dtype列从category重置为object导致内存使用增加，但让我恢复了更重要的运行时间。此外，保留类别列也减慢了我的多进程代码，所以我最终将所有类别列更改为对象。这不是你要寻找的答案，但它证实了你的经验。 - MD_1977

Answer 1

我不确定这个减速是从哪里来的，但一种解决方法是直接存储类别代码：

df3 = df.copy()
animals = pd.Categorical(df['animals'])
df3['animals'] = animals.codes
df3.groupby('days').agg({'animals': 'first'}).apply(lambda code: animals.categories[code])

这并不是最干净的解决方案，因为它需要外部元数据，但它实现了您要寻找的内存效率和计算速度两者。挖掘Pandas在内部执行何种操作导致分类变量变慢会很有趣。

---

编辑：我找到了为什么会发生这种情况......作为first()聚合的一部分，pandas在该列上调用np.asarray()。对于分类列，这将把列转换回非分类列，从而导致不必要的开销。修复此问题将有助于改进Pandas包！

- jakevdp

---

是的，因此，“workaround”、“不是最干净的解决方案”，以及“修复这个问题将是对pandas包有用的贡献”。 - jakevdp