python - 如何使用多索引加速 pandas 的索引？

Question

我有一个多索引的 DataFrame 如下:

In [1]: import pandas as pd
In [2]: import numpy as np
In [3]: symbol = [f'A{i:05d}' for i in range(4000)]
In [4]: date = pd.date_range('20190101', '20201231')
In [5]: index = pd.MultiIndex.from_product([date, symbol], names=['date', 'symbol'])
In [6]: frame = pd.DataFrame(np.random.random((len(index), 4)), index=index, columns=['A', 'B', 'C', 'D'])

我想选择 frame 的子范围，直观的解决方案性能很差:

In [7]: start, end = pd.to_datetime(['20190701', '20190801'])
In [9]: tickers = [f'A{i:05d}' for i in range(4000) if i % 555 != 3]
In [10]: %time a = frame.loc[(slice(start, end), tickers), 'A']
Wall time: 1min 41s

更复杂、更快速的解决方案:

In [11]: %time b = frame['A'].unstack()[tickers].loc[start:end].stack()
Wall time: 616 ms

In [12]: a.equals(b)
Out[12]: True

但是，第二种方案有两个缺点:

1. 它没有第一个优雅；
2. 如果我想选择多个列，例如frame.loc[(slice(start, end), tickers), ['A', 'B']]。

我的问题还有其他快速索引方法吗？

我的python环境:

INSTALLED VERSIONS
------------------
commit           : None
python           : 3.8.3.final.0
python-bits      : 64
OS               : Windows
OS-release       : 10
machine          : AMD64
processor        : Intel64 Family 6 Model 158 Stepping 10, GenuineIntel
byteorder        : little
LC_ALL           : None
LANG             : None
LOCALE           : Chinese (Simplified)_China.936

pandas           : 1.0.5
numpy            : 1.18.5
pytz             : 2020.1
dateutil         : 2.8.1
pip              : 20.1.1
setuptools       : 49.2.0.post20200714
Cython           : None
pytest           : None
hypothesis       : None
sphinx           : 3.1.2
blosc            : None
feather          : None
xlsxwriter       : None
lxml.etree       : None
html5lib         : None
pymysql          : None
psycopg2         : None
jinja2           : 2.11.2
IPython          : 7.16.1
pandas_datareader: None
bs4              : None
bottleneck       : None
fastparquet      : None
gcsfs            : None
lxml.etree       : None
matplotlib       : 3.2.2
numexpr          : 2.7.1
odfpy            : None
openpyxl         : None
pandas_gbq       : None
pyarrow          : None
pytables         : None
pytest           : None
pyxlsb           : None
s3fs             : None
scipy            : 1.5.0
sqlalchemy       : None
tables           : 3.6.1
tabulate         : 0.8.3
xarray           : None
xlrd             : 1.2.0
xlwt             : None
xlsxwriter       : None
numba            : 0.50.1

Answer 1

frame.loc[start:end][frame.loc[start:end].index.isin(tickers, level='symbol')]

这非常快，可以让您获得完整的数据框以选择您想要的任何列，尽管优雅是有争议的(双 loc)