如何改进数据框中的搜索索引

Question

给定一个带有时间戳索引的 pandas 数据帧，已排序。 我有一个标签，我需要找到最接近该标签的索引。 另外，我需要找到一个较小的时间戳，所以搜索应该在次要时间戳中计算。 这是我的代码：

import pandas as pd
import datetime

data = [i for i in range(100)]
dates = pd.date_range(start="01-01-2018", freq="min", periods=100)
dataframe = pd.DataFrame(data, dates)

label = "01-01-2018 00:10:01"
method = "pad"
tol = datetime.timedelta(seconds=60)
idx = dataframe.index.get_loc(key=label, method="pad", tolerance=tol)

print("Closest idx:"+str(idx))
print("Closest date:"+str(dataframe.index[idx]))

搜索太慢了。有什么办法可以改善吗？

Answer 1

为了提高性能，我建议您对正在搜索的内容进行转换。您可以将DateTimeIndex 转换为Unix 时间，而不是使用get_loc，并在底层numpy 数组上使用np.searchsorted（顾名思义，这需要排序索引）。

---

get_loc:

（你目前的做法）

label = "01-01-2018 00:10:01"
tol = datetime.timedelta(seconds=60)
idx = dataframe.index.get_loc(key=label, method="pad", tolerance=tol)
print(dataframe.iloc[idx])

0    10
Name: 2018-01-01 00:10:00, dtype: int64

这是时间安排：

%timeit dataframe.index.get_loc(key=label, method="pad", tolerance=tol)
2.03 ms ± 81.7 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

np.searchsorted:

arr = df.index.astype(int)//10**9
l = pd.to_datetime(label).timestamp()
idx = np.max(np.searchsorted(arr, l, side='left')-1, 0)
print(dataframe.iloc[idx])

0    10
Name: 2018-01-01 00:10:00, dtype: int64

还有时间：

%timeit np.max(np.searchsorted(arr, l, side='left')-1, 0)
56.6 µs ± 979 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

（我没有包括设置成本，因为初始数组创建应该是你做一次，然后用于每个查询，但即使我确实包括了设置成本，这种方法更快）：

%%timeit
arr = df.index.astype(int)//10**9
l = pd.to_datetime(label).timestamp()
np.max(np.searchsorted(arr, l, side='left')-1, 0)

394 µs ± 3.84 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

上述方法不强制执行60s的容差，尽管这很容易检查：

>>> np.abs(arr[idx]-l)<60
True