Pandas中Apply函數加速百倍的技巧分享

2022-07-29 22:02:04

前言

雖然目前dask,cudf等包的出現，使得我們的資料處理大大得到了加速，但是並不是每個人都有比較好的gpu，非常多的朋友仍然還在使用pandas工具包，但有時候真的很無奈，pandas的許多問題我們都需要使用apply函數來進行處理，而apply函數是非常慢的，本文我們就介紹如何加速apply函數600倍的技巧。

實驗對比

01 Apply(Baseline)

我們以Apply為例，原始的Apply函數處理下面這個問題，需要18.4s的時間。

import pandas as pd
import numpy as np
df = pd.DataFrame(np.random.randint(0, 11, size=(1000000, 5)), columns=('a','b','c','d','e'))
def func(a,b,c,d,e):
    if e == 10:
        return c*d
    elif (e < 10) and (e>=5):
        return c+d
    elif e < 5:
        return a+b
%%time
df['new'] = df.apply(lambda x: func(x['a'], x['b'], x['c'], x['d'], x['e']), axis=1)

CPU times: user 17.9 s, sys: 301 ms, total: 18.2 s
Wall time: 18.4 s

02 Swift加速

因為處理是並行的，所以我們可以使用Swift進行加速，在使用Swift之後，相同的操作在我的機器上可以提升到7.67s。

%%time
# !pip install swifter
import swifter
df['new'] = df.swifter.apply(lambda x : func(x['a'],x['b'],x['c'],x['d'],x['e']),axis=1)
HBox(children=(HTML(value='Dask Apply'), FloatProgress(value=0.0, max=16.0), HTML(value='')))

CPU times: user 329 ms, sys: 240 ms, total: 569 ms
Wall time: 7.67 s

03 向量化

使用Pandas和Numpy的最快方法是將函數向量化。如果我們的操作是可以直接向量化的話，那麼我們就儘可能的避免使用：

for迴圈；

列表處理；

apply等操作

在將上面的問題轉化為下面的處理之後，我們的時間縮短為：421 ms。

%%time
df['new'] = df['c'] * df['d'] #default case e = =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']

CPU times: user 134 ms, sys: 149 ms, total: 283 ms
Wall time: 421 ms

04 類別轉化+向量化

我們先將上面的類別轉化為int16型，再進行相同的向量化操作，發現時間縮短為：116 ms。

for col in ('a','b','c','d'):
    df[col] = df[col].astype(np.int16) 
%%time
df['new'] = df['c'] * df['d'] #default case e = =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']

CPU times: user 71.3 ms, sys: 42.5 ms, total: 114 ms
Wall time: 116 ms

05 轉化為values處理

在能轉化為.values的地方儘可能轉化為.values，再進行操作。

此處先轉化為.values等價於轉化為numpy，這樣我們的向量化操作會更加快捷。

於是，上面的操作時間又被縮短為：74.9ms。

%%time
df['new'] = df['c'].values * df['d'].values #default case e = =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']