代码之家 › 专栏 › 技术社区 › danie

如何提高分组/应用效率

pandas-groupby dataframe pandas performance python

danie · 技术社区 · 8 年前

我有一个以年为索引的熊猫数据框架,一列有股票ID,第二列有回报。数据帧大约有20万行。我想再加上3列,分别是未来5年、10年和20年每只股票的累计回报率。为此,我通过ID列进行分组,并将一个函数应用于分组对象,我在下面的一个简单示例中展示了这个函数。我知道这需要一些时间,但到目前为止,代码已经执行了23个小时,仍在运行。

那么我有两个问题:

为什么python要花这么多时间来执行代码?哪里瓶颈是什么?
关于如何更改代码以加快速度,有什么想法吗?

下面是我的代码,应用于一个更简单的示例。

In [1]: import pandas as pd

In [2]: simple_df = pd.DataFrame([[1,1,1,2,2],[0.1,0.05,0.15,0.3,0.2]], columns=[2010,2011,2012,2011,2012], index=['ID','Return']).T

In [3]: simple_df
Out[3]: 
       ID  Return
2010  1.0    0.10
2011  1.0    0.05
2012  1.0    0.15
2011  2.0    0.30
2012  2.0    0.20

In [4]: grouped = simple_df.groupby('ID', sort=False)

In [5]: create_df = lambda x: pd.DataFrame({i: x.Return.shift(-i) for i in range(0,3)})

In [6]: df_1 = grouped.apply(create_df)

In [7]: df_1
Out[7]: 
         0     1     2
2010  0.10  0.05  0.15
2011  0.05  0.15   NaN
2012  0.15   NaN   NaN
2011  0.30  0.20   NaN
2012  0.20   NaN   NaN

In [8]: df_2 =(df_1+1).cumprod(axis=1)-1

In [9]: df_2
Out[9]: 
         0       1        2
2010  0.10  0.1550  0.32825
2011  0.05  0.2075      NaN
2012  0.15     NaN      NaN
2011  0.30  0.5600      NaN
2012  0.20     NaN      NaN

In [10]: simple_df['Return_3y'] = df_2.iloc[:,2]

In [11]: simple_df
Out[11]: 
       ID  Return  Return_3y
2010  1.0    0.10    0.32825
2011  1.0    0.05        NaN
2012  1.0    0.15        NaN
2011  2.0    0.30        NaN
2012  2.0    0.20        NaN

1 回复 | 直到 8 年前

usr 8 年前

而不是 apply 使用 DataFrameGroupBy.shift 具有 concat :

np.random.seed(234)
N = 10000

idx = np.random.randint(1990, 2020, size=N)
simple_df = pd.DataFrame({'ID':np.random.randint(1000, size=N),
              'Return':np.random.rand(N)}, index=idx).sort_values('ID')

print (simple_df)

In [147]: %%timeit
     ...: grouped = simple_df.groupby('ID', sort=False)
     ...: create_df = lambda x: pd.DataFrame({i: x.Return.shift(-i) for i in range(0,3)})
     ...: df_1 = grouped.apply(create_df)
     ...: df_2 =(df_1+1).cumprod(axis=1)-1
     ...: 
1.01 s Â± 6.3 ms per loop (mean Â± std. dev. of 7 runs, 1 loop each)

In [148]: %%timeit
     ...: g = simple_df.groupby('ID', sort=False)
     ...: df2 = pd.concat([g['Return'].shift(-i) for i in range(3)], axis=1, keys=range(3))
     ...: df2 =(df2+1).cumprod(axis=1)-1
     ...: 
3.91 ms Â± 53.3 Âµs per loop (mean Â± std. dev. of 7 runs, 100 loops each)