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 优化