如何加速非常缓慢的Pandas apply函数？

Question

我有一个非常大的pandas数据集，在某个时候我需要使用以下函数

def proc_trader(data):
    data['_seq'] = np.nan
    # make every ending of a roundtrip with its index
    data.ix[data.cumq == 0,'tag'] = np.arange(1, (data.cumq == 0).sum() + 1)
    # backfill the roundtrip index until previous roundtrip;
    # then fill the rest with 0s (roundtrip incomplete for most recent trades)
    data['_seq'] =data['tag'].fillna(method = 'bfill').fillna(0)
    return data['_seq']
    # btw, why on earth this function returns a dataframe instead of the series `data['_seq']`??

我使用apply方法。

reshaped['_spell']=reshaped.groupby(['trader','stock'])[['cumq']].apply(proc_trader)

显然，我无法在此分享数据，但您是否在我的代码中发现了瓶颈？这可能是 arange 的问题吗？数据中有许多 name-productid 组合。

最小化工作示例：

import pandas as pd
import numpy as np

reshaped= pd.DataFrame({'trader' : ['a','a','a','a','a','a','a'],'stock' : ['a','a','a','a','a','a','b'], 'day' :[0,1,2,4,5,10,1],'delta':[10,-10,15,-10,-5,5,0] ,'out': [1,1,2,2,2,0,1]})


reshaped.sort_values(by=['trader', 'stock','day'], inplace=True)
reshaped['cumq']=reshaped.groupby(['trader', 'stock']).delta.transform('cumsum')
reshaped['_spell']=reshaped.groupby(['trader','stock'])[['cumq']].apply(proc_trader).reset_index()['_seq']

- ℕʘʘḆḽḘ

---

你尝试过对代码进行行级分析吗？ - EnricoGiampieri

有一个库可以做到这一点，而且工作得非常好，请查看！我会去看看的，但它可能是数据相关的。https://github.com/rkern/line_profiler - EnricoGiampieri

你的代码给了我一个 TypeError: incompatible index of inserted column with frame index 错误。有什么遗漏吗？ - DSM

同样出错，TypeError。 - EnricoGiampieri

哈哈，好的技巧，PANDAS！谢谢伙计。 - ℕʘʘḆḽḘ

Answer 1

这里没有什么特别的，只是在一些地方进行了微调。没有必要编写一个函数，所以我没有写。在这个小样本数据上，它比原始代码快了大约两倍。

reshaped.sort_values(by=['trader', 'stock','day'], inplace=True)
reshaped['cumq']=reshaped.groupby(['trader', 'stock']).delta.cumsum()
reshaped.loc[ reshaped.cumq == 0, '_spell' ] = 1
reshaped['_spell'] = reshaped.groupby(['trader','stock'])['_spell'].cumsum()
reshaped['_spell'] = reshaped.groupby(['trader','stock'])['_spell'].bfill().fillna(0)

结果：

   day  delta  out stock trader  cumq  _spell
0    0     10    1     a      a    10     1.0
1    1    -10    1     a      a     0     1.0
2    2     15    2     a      a    15     2.0
3    4    -10    2     a      a     5     2.0
4    5     -5    2     a      a     0     2.0
5   10      5    0     a      a     5     0.0
6    1      0    1     b      a     0     1.0

- JohnE