尋找在巨大的 Pandas Dataframe 中切片一行的最快方法

Question

我的程序需要根據來自巨大 Pandas 數據幀的列中的值獲取一行。 響應時間很關鍵。 我使用最常見的方法來做到這一點，例如：

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

每timeit我的Mac上花了4毫秒來完成對數據框上面的操作以1萬行。 但我的目標是將時間減少到 0.4 毫秒。 我曾經考慮將此數據幀轉換為 Set 但 Set 沒有排序，並且本身不支持索引或切片。 有什么建議？

Answer 1

讓我們設置一下：

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

然后讓我們對一些選項進行基准測試。 您當前的 boolean + .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

即使您必須將其重新轉換為某個系列，這也比以前快了一個數量級。 （如果需要，您也可以很容易地將一個系列映射回 dict 並保持 dict 查找的速度而沒有開銷）

Answer 2

檢查df.query('id == 500000') 的工作速度。