Pandas tshift 在分组中慢

Question

使用Pandas的tshift功能真的很棒，速度也很快！

df = pd.DataFrame(index=pd.date_range(pd.datetime(1970,1,1),pd.datetime(1970,2,1)))
df['data']=.5
%timeit df.sum()
#10000 loops, best of 3: 162 µs per loop 
%timeit df.tshift(-1)
#1000 loops, best of 3: 307 µs per loop #x2 slower

但是，当我在进行groupby操作后再使用tshift时，速度就慢很多了：

df = pd.DataFrame(index=pd.date_range(pd.datetime(1970,1,1),pd.datetime(1970,2,1)))
df['data']=.5
df['A'] = randint(0,2,len(df.index))
%timeit df.groupby('A').sum()
#100 loops, best of 3: 2.72 ms per loop
%timeit df.groupby('A').tshift(-1)
#10 loops, best of 3: 16 ms per loop #x6 slower!

为什么在分组后使用tshift会变得这么慢呢？有没有什么方法可以让它更快？

更新：

我实际的使用场景更接近下面的代码。我发现慢下来的程度和分组的数量有关。

n_A = 50
n_B = 5
index = pd.MultiIndex.from_product([arange(n_A),
                                     arange(n_B),
                                     pd.date_range(pd.datetime(1975,1,1),
                                                   pd.datetime(2010,1,1),
                                                   freq='5AS')],
                                   names=['A', 'B', 'Year'])

df = pd.DataFrame(index=index)
df['data']=.5

%timeit df.reset_index(['A','B']).groupby(['A','B']).sum()
#100 loops, best of 3: 4.34 ms per loop
%timeit df.reset_index(['A','B']).groupby(['A','B']).tshift(-1, freq='5AS')
#10 loops, best of 3: 198 ms per loop # X44 slowdown.

如果我们增加A和B组的数量：

n_A = 500
n_B = 50
...
%timeit df.reset_index(['A','B']).groupby(['A','B']).sum()
#10 loops, best of 3: 35.8 ms per loop
%timeit df.reset_index(['A','B']).groupby(['A','B']).tshift(-1, freq='5AS')
#1 loops, best of 3: 20.3 s per loop # X567 slowdown

我很惊讶，慢下来的程度随着分组数量的增加而增加！有没有更聪明的方法来处理这个问题？

Answer 1

tshift 这个函数在使用时需要一个频率参数（因为在分组后，频率通常是不规律的），所以当你写 df.groupby('A').tshift(-1) 时，它会返回一个空的框架（因为它在每个分组中都在抛出错误，这样会导致速度变慢）。

In [44]: %timeit df.groupby('A').tshift(-1,'D')
100 loops, best of 3: 3.57 ms per loop

In [45]: %timeit df.groupby('A').sum()
1000 loops, best of 3: 1.02 ms per loop

除此之外，这个问题 在这里 也在等待一个用 Cython 实现的 shift（和 tshift）功能。这样的话，它的性能就能和已经用 Cython 实现的 sum 相当了。欢迎大家贡献代码！

使用你的第二个数据集（更大的分组），你可以这样做：

In [59]: def f(df):
   ....:     x = df.reset_index()
   ....:     x['Year_ts'] = pd.DatetimeIndex(x['Year'])-pd.offsets.YearBegin(5)
   ....:     return x.drop(['Year'],axis=1).rename(columns={'Year_ts' : 'Year'}).set_index(['A','B','Year'])
   ....: 

In [60]: result = df.reset_index(['A','B']).groupby(['A','B']).tshift(-1,'5AS')

In [61]: %timeit df.reset_index(['A','B']).groupby(['A','B']).tshift(-1,'5AS')
1 loops, best of 3: 10.8 s per loop

In [62]: result2 = f(df)

In [63]: %timeit f(df)
1 loops, best of 3: 2.51 s per loop

In [64]: result.equals(result2)
Out[64]: True

所以把日期的减法放在分组外面做，这样速度大约快了 4 倍。而且这样做（还有缓存）可能是让分组的 tshift 更快的第一步。