python - pandas中groupby处理 `category`类型时的性能

Question

我想了解 pandas 中使用 category 背后的微妙之处。

我通过创建了一个随机三列 DataFrame

import pandas as pd
import numpy as np

a = np.random.choice(['a', 'A'], size=100000)
b = np.random.choice(range(10000), size=100000)

df = pd.DataFrame(data=dict(A1=a, A2=a, B=b))
df['A2'] = df.A2.astype('category')

处理A2有一个category类型列。

df.info()

<class 'pandas.core.frame.DataFrame'>
RangeIndex: 100000 entries, 0 to 99999
Data columns (total 3 columns):
A1    100000 non-null object
A2    100000 non-null category
B     100000 non-null int64
dtypes: category(1), int64(1), object(1)
memory usage: 1.6+ MB

在 B 上运行 groupby 运算符并在 A1 和 A2 上应用一个简单的函数会导致性能方面的巨大差异，

%%timeit
df_ = df.groupby(by='B').agg(
    dict(
        A1=lambda s: len(s.unique())
    )
)

1 loop, best of 3: 666 ms per loop

和

%%timeit
df_ = df.groupby(by='B').agg(
    dict(
        A2=lambda s: len(s.unique())
    )
)

1 loop, best of 3: 2.73 s per loop

您能告诉我这背后的原因吗？

Answer 1

我认为你可以使用SeriesGroupBy.nunique ，它与 category 一起工作只是稍微慢一点:

print (df.groupby(by='B')['A1'].nunique())
print (df.groupby(by='B')['A2'].nunique())

In [71]: %timeit (df.groupby(by='B')['A1'].nunique())
10 loops, best of 3: 19.8 ms per loop

In [72]: %timeit (df.groupby(by='B')['A2'].nunique())
10 loops, best of 3: 20.2 ms per loop

有趣，如果将 agg 与 dict 一起使用，性能是相同的:

In [77]: %timeit df.groupby(by='B').agg({'A1':'nunique'})
100 loops, best of 3: 20.1 ms per loop

In [78]: %timeit df.groupby(by='B').agg({'A2':'nunique'})
10 loops, best of 3: 20.1 ms per loop

但是如果使用lambda函数，dtypecategory会更慢(与您相同的问题):

In [73]: %timeit df.groupby(by='B').agg(dict(A1=lambda s: s.nunique()))
1 loop, best of 3: 824 ms per loop

In [74]: %timeit df.groupby(by='B').agg(dict(A2=lambda s: s.nunique()))
1 loop, best of 3: 3.07 s per loop