快速选择Python中层次索引pandas数据的时间戳范围
如果你有一个带有时区信息的时间索引的DataFrame，下面的方法可以快速选择两个日期之间的多行数据，左边的日期包含在内，右边的日期不包含：
import pandas as pd
start_ts = pd.Timestamp('20000101 12:00 UTC')
end_ts = pd.Timestamp('20000102 12:00 UTC')
ix_df = pd.DataFrame(0, index=[pd.Timestamp('20000101 00:00 UTC'), pd.Timestamp('20000102 00:00 UTC')], columns=['a'])
EPSILON_TIME = pd.tseries.offsets.Nano()
ix_df[start_ts:end_ts-EPSILON_TIME]
这个方法效率还不错，因为我们没有创建一个临时的索引迭代器，也没有在Python中运行一个lambda表达式来生成新的数据框。实际上，我认为这个选择的复杂度大约是O(log(N))，最多也就这样。我在想，这种选择方法是否也可以在MultiIndex的某个特定层级上使用，还是说我必须创建一个临时的迭代器或者运行lambda表达式。例如：
mux = pd.MultiIndex.from_arrays([[pd.Timestamp('20000102 00:00 UTC'), pd.Timestamp('20000103 00:00 UTC')], [pd.Timestamp('20000101 00:00 UTC'), pd.Timestamp('20000102 00:00 UTC')]])
mux_df = pd.DataFrame(0, index=mux, columns=['a'])
然后我可以用同样的方法在索引的第一个（零级）层级上进行选择：
mux_df[start_ts:end_ts-EPSILON_TIME]
这样得到的结果是：
a
2000-01-02 00:00:00+00:00 2000-01-01 00:00:00+00:00 0
但是对于第二个层级，我就得选择一个比较慢的方法：
values_itr = mux_df.index.get_level_values(1)
mask_ser = (values_itr >= start_ts) & (values_itr < end_ts)
mux_df[mask_ser]
这样得到的结果是正确的：
a
2000-01-03 00:00:00+00:00 2000-01-02 00:00:00+00:00 0
有没有什么快速的解决办法？谢谢！
编辑：选择的方法
最后我选择了这个方法，因为我意识到我还需要切片：
def view(data_df):
if len(data_df.index) == 0:
return data_df
values_itr = data_df.index.get_level_values(0)
values_itr = values_itr.values
from_i = np.searchsorted(values_itr, np.datetime64(start_ts), side='left')
to_i = np.searchsorted(values_itr, np.datetime64(end_ts), side='left')
return data_df.ix[from_i:to_i]
然后执行view(data_df).copy()。注意：我在索引的第一层级中的值实际上是已经排序的。
1 个回答
3
其实你在这里比较的就像是苹果和橘子。
In [59]: N = 1000000
In [60]: pd.set_option('max_rows',10)
In [61]: idx = pd.IndexSlice
In [62]: df = DataFrame(np.arange(N).reshape(-1,1),columns=['value'],index=pd.MultiIndex.from_product([list('abcdefghij'),date_range('20010101',periods=N/10,freq='T',tz='US/Eastern')],names=['one','two']))
In [63]: df
Out[63]:
value
one two
a 2001-01-01 00:00:00-05:00 0
2001-01-01 00:01:00-05:00 1
2001-01-01 00:02:00-05:00 2
2001-01-01 00:03:00-05:00 3
2001-01-01 00:04:00-05:00 4
... ...
j 2001-03-11 10:35:00-05:00 999995
2001-03-11 10:36:00-05:00 999996
2001-03-11 10:37:00-05:00 999997
2001-03-11 10:38:00-05:00 999998
2001-03-11 10:39:00-05:00 999999
[1000000 rows x 1 columns]
In [64]: df2 = df.reset_index(level='one').sort_index()
df
In [65]: df2
Out[65]:
one value
two
2001-01-01 00:00:00-05:00 a 0
2001-01-01 00:00:00-05:00 i 800000
2001-01-01 00:00:00-05:00 h 700000
2001-01-01 00:00:00-05:00 g 600000
2001-01-01 00:00:00-05:00 f 500000
... .. ...
2001-03-11 10:39:00-05:00 c 299999
2001-03-11 10:39:00-05:00 b 199999
2001-03-11 10:39:00-05:00 a 99999
2001-03-11 10:39:00-05:00 i 899999
2001-03-11 10:39:00-05:00 j 999999
[1000000 rows x 2 columns]
当我重置索引（也就是创建一个单层索引）时，它就不再是唯一的了。这一点很重要，因为它的搜索方式会不同。所以你不能简单地把单层唯一索引和多层索引的性能进行比较。
结果发现，使用多层索引切片器（在0.14.0版本中引入的）可以让任何层级的索引速度都很快。
In [66]: %timeit df.loc[idx[:,'20010201':'20010301'],:]
1 loops, best of 3: 188 ms per loop
In [67]: df.loc[idx[:,'20010201':'20010301'],:]
Out[67]:
value
one two
a 2001-02-01 00:00:00-05:00 44640
2001-02-01 00:01:00-05:00 44641
2001-02-01 00:02:00-05:00 44642
2001-02-01 00:03:00-05:00 44643
2001-02-01 00:04:00-05:00 44644
... ...
j 2001-03-01 23:55:00-05:00 986395
2001-03-01 23:56:00-05:00 986396
2001-03-01 23:57:00-05:00 986397
2001-03-01 23:58:00-05:00 986398
2001-03-01 23:59:00-05:00 986399
[417600 rows x 1 columns]
把这个和一个非唯一的单层索引比较一下
In [68]: %timeit df2.loc['20010201':'20010301']
1 loops, best of 3: 470 ms per loop
这里是一个唯一的单层索引
In [73]: df3 = DataFrame(np.arange(N).reshape(-1,1),columns=['value'],index=date_range('20010101',periods=N,freq='T',tz='US/Eastern'))
In [74]: df3
Out[74]:
value
2001-01-01 00:00:00-05:00 0
2001-01-01 00:01:00-05:00 1
2001-01-01 00:02:00-05:00 2
2001-01-01 00:03:00-05:00 3
2001-01-01 00:04:00-05:00 4
... ...
2002-11-26 10:35:00-05:00 999995
2002-11-26 10:36:00-05:00 999996
2002-11-26 10:37:00-05:00 999997
2002-11-26 10:38:00-05:00 999998
2002-11-26 10:39:00-05:00 999999
[1000000 rows x 1 columns]
In [75]: df3.loc['20010201':'20010301']
Out[75]:
value
2001-02-01 00:00:00-05:00 44640
2001-02-01 00:01:00-05:00 44641
2001-02-01 00:02:00-05:00 44642
2001-02-01 00:03:00-05:00 44643
2001-02-01 00:04:00-05:00 44644
... ...
2001-03-01 23:55:00-05:00 86395
2001-03-01 23:56:00-05:00 86396
2001-03-01 23:57:00-05:00 86397
2001-03-01 23:58:00-05:00 86398
2001-03-01 23:59:00-05:00 86399
[41760 rows x 1 columns]
到目前为止，速度是最快的
In [76]: %timeit df3.loc['20010201':'20010301']
1 loops, best of 3: 294 ms per loop
最好的是没有时区的单层唯一索引
In [77]: df3 = DataFrame(np.arange(N).reshape(-1,1),columns=['value'],index=date_range('20010101',periods=N,freq='T'))
In [78]: %timeit df3.loc['20010201':'20010301']
1 loops, best of 3: 240 ms per loop
而且这是迄今为止最快的方法（我这里做了一个稍微不同的搜索，以获得相同的结果，因为上面搜索的语义包括了指定日期的所有日期）
In [101]: df4 = df3.reset_index()
In [103]: %timeit df4.loc[(df4['index']>='20010201') & (df4['index']<'20010302')]
100 loops, best of 3: 10.6 ms per loop
In [104]: df4.loc[(df4['index']>='20010201') & (df4['index']<'20010302')]
Out[104]:
index value
44640 2001-02-01 00:00:00 44640
44641 2001-02-01 00:01:00 44641
44642 2001-02-01 00:02:00 44642
44643 2001-02-01 00:03:00 44643
44644 2001-02-01 00:04:00 44644
... ... ...
86395 2001-03-01 23:55:00 86395
86396 2001-03-01 23:56:00 86396
86397 2001-03-01 23:57:00 86397
86398 2001-03-01 23:58:00 86398
86399 2001-03-01 23:59:00 86399
[41760 rows x 2 columns]
那么，为什么第四种方法是最快的呢？它构建了一个布尔索引数组，然后使用非零值，所以速度非常快。前面三种方法在确定索引是唯一且单调的情况下，使用了searchsorted（两次）来找出端点，所以会有多个步骤在进行。
总之，布尔索引的速度非常快，所以要多用它！不过结果可能会不同，前面三种方法的速度也可能会根据你选择的内容而有所不同，比如选择的范围小可能会有不同的性能表现。