如何加速apply函数600倍的技巧

↓推荐关注↓

[ 引言 ] 虽然目前dask,cudf等包的出现，使得我们的数据处理大大得到了加速，但是并不是每个人都有比较好的gpu，非常多的朋友仍然还在使用pandas工具包，但有时候真的很无奈，pandas的许多问题我们都需要使用apply函数来进行处理，而apply函数是非常慢的，本文我们就介绍如何加速apply函数600倍的技巧。

实验对比01Apply(Baseline) 我们以Apply为例，原始的Apply函数处理下面这个问题，需要18.4s的时间。

importpandasaspd
importnumpyasnp
df=pd.DataFrame(np.random.randint(0,11,size=(1000000,5)),columns=('a','b','c','d','e'))
deffunc(a,b,c,d,e):
ife==10:
returnc*d
elif(e< 10) and (e>=5):
returnc+d
elife< 5:
        returna+b
%%time
df['new']=df.apply(lambdax:func(x['a'],x['b'],x['c'],x['d'],x['e']),axis=1)
CPUtimes:user17.9s,sys:301ms,total:18.2s
Walltime:18.4s

02Swift加速 因为处理是并行的，所以我们可以使用Swift进行加速，在使用Swift之后，相同的操作在我的机器上可以提升到7.67s。

%%time
#!pipinstallswifter
importswifter
df['new']=df.swifter.apply(lambdax:func(x['a'],x['b'],x['c'],x['d'],x['e']),axis=1)
HBox(children=(HTML(value='DaskApply'),FloatProgress(value=0.0,max=16.0),HTML(value='')))

CPUtimes:user329ms,sys:240ms,total:569ms
Walltime:7.67s

03向量化 使用Pandas和Numpy的最快方法是将函数向量化。如果我们的操作是可以直接向量化的话，那么我们就尽可能的避免使用：

• for循环；
• 列表处理；
• apply等操作

在将上面的问题转化为下面的处理之后，我们的时间缩短为：421 ms。

%%time
df['new']=df['c']*df['d']#defaultcasee==10
mask=df['e']< 10
df.loc[mask,'new']=df['c']+df['d']
mask=df['e']< 5
df.loc[mask,'new']=df['a']+df['b']
CPUtimes:user134ms,sys:149ms,total:283ms
Walltime:421ms

04类别转化+向量化 我们先将上面的类别转化为int16型，再进行相同的向量化操作，发现时间缩短为：116 ms。

forcolin('a','b','c','d'):
df[col]=df[col].astype(np.int16)
%%time
df['new']=df['c']*df['d']#defaultcasee==10
mask=df['e']< 10
df.loc[mask,'new']=df['c']+df['d']
mask=df['e']< 5
df.loc[mask,'new']=df['a']+df['b']
CPUtimes:user71.3ms,sys:42.5ms,total:114ms
Walltime:116ms

05转化为values处理 在能转化为.values的地方尽可能转化为.values，再进行操作。

• 此处先转化为.values等价于转化为numpy，这样我们的向量化操作会更加快捷。

于是，上面的操作时间又被缩短为：74.9ms。

%%time
df['new']=df['c'].values*df['d'].values#defaultcasee==10
mask=df['e'].values< 10
df.loc[mask,'new']=df['c']+df['d']
mask=df['e'].values< 5
df.loc[mask,'new']=df['a']+df['b']
CPUtimes:user64.5ms,sys:12.5ms,total:77ms
Walltime:74.9ms

实验汇总 通过上面的一些小的技巧，我们将简单的Apply函数加速了几百倍，具体的：

• Apply: 18.4 s
• Apply + Swifter: 7.67 s
• Pandas vectorizatoin: 421 ms
• Pandas vectorization + data types: 116 ms
• Pandas vectorization + values + data types: 74.9ms

参考文献：Do You Use Apply in Pandas? There is a 600x Faster Way

审核编辑 ：李倩