按类别过滤熊猫数据框的最快方法

Question

我有一个非常大的数据框，其中包含 1 亿行和分类列。我想知道是否有比使用.isin() 方法或.join() 提到的here 方法更快的按类别选择行的方法。

鉴于数据已经分类，我认为选择类别应该很快，但我运行的一些测试表现令人失望。我发现的唯一其他解决方案来自 here，但该解决方案似乎不适用于 pandas 0.20.2。

这是一个示例数据集。

import pandas as pd
import random
import string
df = pd.DataFrame({'categories': [random.choice(string.ascii_letters) 
                                  for _ in range(1000000)]*100,
                   'values': [random.choice([0,1]) 
                              for _ in range(1000000)]*100})
df['categories'] = df['categories'].astype('category')

使用.isin()进行测试：

%timeit df[df['categories'].isin(list(string.ascii_lowercase))]
44 s ± 894 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

使用.join()：

%timeit df.set_index('categories').join(
    pd.Series(index=list(string.ascii_lowercase), name='temp'), 
    how='inner').rename_axis('categories').reset_index().drop('temp', 1)
24.7 s ± 1.69 s per loop (mean ± std. dev. of 7 runs, 1 loop each)

Answer 1

这是一种类似但不同的方法，它直接比较值而不是使用isin。

基本map/lambda 比较：

%timeit df[df['categories'].map(lambda x: x in string.ascii_lowercase)]
> 1 loop, best of 3: 12.3 s per loop

使用isin：

%timeit df[df['categories'].isin(list(string.ascii_lowercase))]
> 1 loop, best of 3: 55.1 s per loop

版本：Py 3.5.1 / IPython 5.1.0 / Pandas 0.20.3

背景：我在one of the SO posts 中注意到您链接到的评论者提到isin 需要在执行期间创建set()，因此跳过该步骤并进行基本的列表查找似乎是这里的加速。

Disclamer：不是我经常处理的秤类型，因此可能会有更快的选择。

编辑：JohnGalt 在 cmets 中请求的额外细节：

df.shape
> (100000000, 2)
df.dtypes
> categories    category
 values           int64
 dtype: object

为了创建示例数据，我从最初的问题中复制/粘贴了示例 DF。在 MBP 上运行，2015 年初的模型。

Answer 2

这是对 1/10 样本的运行。 Groupby 自然地利用了分类。目前，为 Categorical 实现 df.loc[list_of_items] 效率不是很高，但可以在后台利用它。

In [68]: %timeit -n 3 -r 1 df[df['categories'].isin(list(string.ascii_lowercase))]
3.67 s +- 0 ns per loop (mean +- std. dev. of 1 run, 3 loops each)

In [69]: s = set(list(string.ascii_lowercase))

In [70]: %timeit -n 3 -r 1 df.groupby('categories', as_index=False).filter(lambda x: x.name in s)
1.09 s +- 0 ns per loop (mean +- std. dev. of 1 run, 3 loops each)

In [71]: result2 = df.groupby('categories', as_index=False).filter(lambda x: x.name in s)

In [72]: result1 = df[df['categories'].isin(list(string.ascii_lowercase))]
result1.equals

In [73]: result1.equals(result2)
Out[73]: True