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

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

为提高性能，我建议对您要搜索的内容进行转换。您可以将get_loc转换为Unix Time，而不是使用DateTimeIndex，并在基础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