写高性能的Pandas代码

写高性能的Pandas代码

我觉得吧，python作为科学计算的最常使用语言之一，应对大量的数据计算，如果太慢了，会让需要不断试错的科学计算方法消耗过多的时间。所以我常常在思考，python到底有多慢，让大家一开始用就觉得它慢？又有多快，让大家都用它来进行上GB数据的计算？

pandas是用来处理科学计算数据的最常用框架，pandas的性能怎么样呢？在一步步尝试中，我发现这取决于代码的写法。接下来就来比较一下，遍历数据集这种情景下几种写法的性能消耗是怎样的。

数据采用的是网上随便找的一个数据集：

import pandas as pd
import numpy as np

data = pd.read_csv("https://vincentarelbundock.github.io/Rdatasets/csv/datasets/EuStockMarkets.csv")
data.head()

	Unnamed: 0	DAX	SMI	CAC	FTSE
0	1	1628.75	1678.1	1772.8	2443.6
1	2	1613.63	1688.5	1750.5	2460.2
2	3	1606.51	1678.6	1718.0	2448.2
3	4	1621.04	1684.1	1708.1	2470.4
4	5	1618.16	1686.6	1723.1	2484.7

data.describe()

	Unnamed: 0	DAX	SMI	CAC	FTSE
count	1860.000000	1860.000000	1860.000000	1860.000000	1860.000000
mean	930.500000	2530.656882	3376.223710	2227.828495	3565.643172
std	537.080069	1084.792740	1663.026465	580.314198	976.715540
min	1.000000	1402.340000	1587.400000	1611.000000	2281.000000
25%	465.750000	1744.102500	2165.625000	1875.150000	2843.150000
50%	930.500000	2140.565000	2796.350000	1992.300000	3246.600000
75%	1395.250000	2722.367500	3812.425000	2274.350000	3993.575000
max	1860.000000	6186.090000	8412.000000	4388.500000	6179.000000

就用DAX这列，来进行测试，用来测试的函数大概就是 sin(DAX) * 1.1，并没有什么特殊的意义，就是纯粹消耗时间。

大概我们会测试下面的几种写法：

1. 朴素for循环
2. iterrows方法循环
3. apply方法
4. 向量方法

来一一测试消耗的时间，统计时间方法统一为timeit。

先来看第一种方法吧

朴素for循环

%%timeit

result = [np.sin(data.iloc[i]['DAX']) * 1.1 for i in range(len(data))]

data['target'] = result

422 ms ± 29.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

使用普通的for循环，1860次循环，耗时400+ms，这个数字在每台电脑上运行应该都不同，但是和下面几个写法的运行时间肯定是相对的。

这种写法的时间消耗很大，很python。这种写法也是很不推荐的，所以来改善一下

iterrows方法循环

pandas提供了iterrows方法，提供了类似enumerate的行为，可以拿到index和当前循环的对象。

iterrows的性能较朴素的for循环来说，性能提高了不少，使用起来也很方便。同时，pandas也提供了一个itertuples方法，能提供很高的性能，但是牺牲了方便程度，不仅没有了循环的index，而且循环的对象如其名所说，变成了一个tuple，所以取值不能通过其列的index名称来获取，只能通过tuple的index来取值。

%%timeit

result = [np.sin(row['DAX']) * 1.1 for index,row in data.iterrows()]

data['target'] = result

103 ms ± 1.15 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

%%timeit

result = [np.sin(row[2]) * 1.1 for row in data.itertuples()]

data['target'] = result

9.67 ms ± 649 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

除了iterrows这种方法之外，还有另外的一种方法。

apply方法

这个方法以回调方法的形式来遍历整个DataFrame，同时可以使用axis参数来指定遍历的轴（以行遍历还是以列遍历）。

apply方法的性能好于iterrows但是弱于itertuples，方便程度和iterrows不相上下。

%%timeit

data['target'] = data.apply(lambda row: np.sin(row['DAX']) * 1.1, axis=1)

60.7 ms ± 3.14 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

在使用itertuples的情况下，性能最高已经优化到了9ms左右，大约优化了10倍，那么还能继续加强吗

向量方法

如果进行科学计算，应该知道向量。DataFrame中的数据基本上都是向量，包括筛选条件。

所以如果用DataFrame的向量单位来进行计算，是否要更快呢？

%%timeit

data['target'] = np.sin(data['DAX']) * 1.1

427 µs ± 36.4 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

可以看到，这样写更加简洁，效率也高的吓人，只需要400+us，比最初的400ms优化了1000倍，已经小于毫秒了。

思考

如果深入思考，应该能明白为什么能一步步的优化到现在的性能。

朴素的for循环携带的信息太多并且大多都是不需要使用的，而且使用的是python缓慢的循环。那么可以针对这两点来进行优化，首先优化掉python循环，使用apply函数，使用C代码来循环，这样性能就能提高一倍了；然后减少携带信息，只使用不可变的tuple，能够提高相当多的性能。

但其实tuple并不是DataFrame的原有结构，转换成tuple还是需要花费很多时间，能直接使用DataFrame的结构来计算，应该能提高很多性能。

所以直接使用向量来计算，将性能又提高了20倍。

那么，能按照这个思路继续思考，其实DataFrame的向量也带有了一些附加信息，如DataFrame的index，那么舍弃了这些信息，用最底层的计算单位来做运算，应该也能提高一部分性能。

Pandas是使用的Numpy做的基本数据单位，所以能从向量中提取出原来的Numpy数组再进行计算，就能提高一部分性能了。

%%timeit

data['target'] = np.sin(data['DAX'].values) * 1.1

218 µs ± 6.34 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

性能又提高了将近两倍，较最开始提高了2000倍。

总结

在很多系统设计之初，就有很开始担心性能问题，我觉得是完全没有必要的。

系统设计的难点在于用最好或者更好的方式实现这个功能，系统的性能不是系统的瓶颈。Python易于试错，能快速迭代，也有很多性能优化的手段，如果这些手段并不能满足需求，还能使用Cython或者直接使用C来编写部分代码。

当然，性能从来都是Python被大家吐槽的地方，一开始就打上了慢的标签。所以也不需要太过在意性能，用几十行代码做出功能来进行实践，再重写成几千行的C代码也没有什么不可取的嘛。