首页 > 解决方案 > 寻找在一个巨大的 Pandas Dataframe 中分割一行的最快方法

问题描述

我的程序需要根据一个巨大的 Pandas 数据框中的列中的值来获取一行。响应时间至关重要。我使用最常见的方法来做到这一点,例如:

df.loc[df['id'] == 500000, :]

在我的Mactimeit上,在 100 万行的数据帧上完成上述操作需要 4 毫秒。但我的目标是将时间减少到 0.4 毫秒。我曾经考虑将此数据框转换为 Set 但 Set 没有排序,并且本身不支持索引或切片。有什么建议么?

标签: pythonpandas

解决方案


让我们设置一下:

import pandas as pd
import numpy as np
df = pd.DataFrame({"id": np.random.randint(100,size=(1000000,))})

然后让我们对一些选项进行基准测试。您当前的布尔值 + .loc

>>> timeit.timeit("df.loc[df['id'] == 50, :]", setup = "from __main__ import df", number=1000)
2.566220869999597

查询引擎:

>>> timeit.timeit("df.query('id == 50')", setup = "from __main__ import df", number=1000)
14.591400260000228

使用索引作为单独的查找:

>>> idx = pd.Index(df['id'])
>>> timeit.timeit("df.loc[idx == 50, :]", setup = "from __main__ import df, idx", number=1000)
2.2155187300013495

使用数据帧索引进行查找:

>>> df.index = df["id"]
>>> timeit.timeit("df.loc[50, :]", setup = "from __main__ import df", number=1000)
2.625610274999417

评论中的某个人的.isin()想法是:

>>> timeit.timeit("df.loc[df['id'].isin([50]), :]", setup = "from __main__ import df", number=1000)
9.542700138999862

看起来除了查询引擎对于一个简单的相等性很慢(如预期的那样)之外,你不会比你所拥有的查找时间更好。

df_unique = pd.DataFrame({'id': range(1000000)})

让我们看看唯一 ID 有什么帮助:

>>> timeit.timeit("df_unique.loc[df_unique['id'] == 50, :]", setup = "from __main__ import df_unique", number=1000)
1.9672015519990964

然后到一个字典:

>>> df_unique.index = df_unique['id']
>>> df_dict = df_unique.to_dict(orient='index')
>>> timeit.timeit("df_dict[50]", setup = "from __main__ import df_dict", number=1000)
6.247700002859347e-05

好吧,看起来这是一个明显的赢家。

>>> timeit.timeit("pd.Series(df_dict[50])", setup = "from __main__ import df_dict, pd", number=1000)
0.2747819870000967

即使您必须将其重新转换为系列,这也比以前快了一个数量级。(如果需要,您也可以很容易地将一系列映射回字典,并保持字典查找的速度而没有开销)


推荐阅读