Pandas 性能優化 學習筆記

摘要

本文介紹了使用 Pandas 進行數據挖掘時經常使用的加速技巧。

實驗環境

import numpy as np
import pandas as pd
print(np.__version__)
print(pd.__version__)

1.16.5
0.25.2

性能分析工具

本文使用到的性能分析工具，參考：Python 性能評估 學習筆記

數據準備

tsdf = pd.DataFrame(np.random.randint(1, 1000, (1000, 3)), columns=['A', 'B', 'C'],
                    index=pd.date_range('1/1/1900', periods=1000))
tsdf['D'] = np.random.randint(1, 3, (1000, ))
tsdf.head(3)

A   B   C
1900-01-01  820 827 884 1
1900-01-02  943 196 513 1
1900-01-03  693 194 6   2

使用 numpy 數組加速運算

map, applymap, apply 之間的區別，參考：Difference between map, applymap and apply methods in Pandas

apply(func, raw=True)

Finally, apply() takes an argument raw which is False by default, which converts each row or column into a Series before applying the function. When set to True, the passed function will instead receive an ndarray object, which has positive performance implications if you do not need the indexing functionality.
Pandas 官方文檔

DataFrame.apply() 支持參數 raw，爲 True 時，直接將 ndarray 輸入函數，利用 numpy 並行化加速。
有多快？

%%timeit
tsdf.apply(np.mean)  # raw=False (default)
# 740 µs ± 28.5 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%%timeit
tsdf.apply(np.mean, raw=True)
# 115 µs ± 2.76 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

由 740 微秒下降到 115 微秒。
什麼條件下能夠使用？

1. 只有 DataFrame.apply() 支持，Series.apply() 和 Series.map() 均不支持；
2. func 不使用 Series 索引時。

tsdf.apply(np.argmax)  # raw=False, 保留索引

A   2019-12-08
B   2021-03-14
C   2020-04-09
D   2019-11-30
dtype: datetime64[ns]

tsdf.apply(np.argmax, raw=True)  # 索引丟失

A      8
B    470
C    131
D      0
dtype: int64

.values

多個 Series 計算時，能夠使用 .values 將 Series 轉換爲 ndarray 再計算。

%%timeit
tsdf.A * tsdf.B
# 123 µs ± 2.86 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%%timeit
tsdf.A.values * tsdf.B.values
# 11.1 µs ± 1.09 µs per loop (mean ± std. dev. of 7 runs, 100000 loops each)

由 123 微秒下降到 11 微秒。
補充說明
注意到 Pandas 0.24.0 引入了 .array 和 .to_numpy()，參考。但這兩種方法的速度不如 values，建議在數據爲數值的狀況下繼續使用 values。

%%timeit
tsdf.A.array * tsdf.B.array
# 37.9 µs ± 938 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%%timeit
tsdf.A.to_numpy() * tsdf.B.to_numpy()
# 15.6 µs ± 110 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

可見兩種方法均慢於 values 的 11 微秒。

字符串操做優化

數據準備

tsdf['S'] = tsdf.D.map({1: '123_abc', 2: 'abc_123'})

%%timeit
tsdf.S.str.split('_', expand=True)[0]  # 獲得'_'以前的字符串
# 1.44 ms ± 97.5 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

一種優化思路是：針對特定場景，不須要使用 split，能夠改用 partition：

%%timeit
tsdf.S.str.partition('_', expand=True)[0]
# 1.39 ms ± 44.8 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

速度略有提高。試試 apply ：

%%timeit
tsdf.S.apply(lambda a: a.partition('_')[0])
# 372 µs ± 8.1 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

可見使用 apply 速度反而比 Pandas 自帶的字符串處理方法要快，這多是由於 Pandas 支持的數據類型多，處理過程當中存在一些冗餘的判斷。
注意到原有數據只有2種，理論上對每一種數據取值只須要計算一次，其它值直接 map 就行。所以考慮轉換爲 Categorical 類型：

tsdf['S_category'] = tsdf.S.astype('category')

%%timeit
tsdf.S_category.apply(lambda a: a.partition('_')[0])
# 246 µs ± 3.36 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

耗時下降至 246 微秒。

IO 優化