最近做科研时经常需要遍历整个 DataFrame，进行各种列操作，例如把某列的值全部转成 pd.Timestamp 格式或者将某两列的值进行 element-wise 运算之类的。大数据的数据量随便都是百万条起跳，如果只用 for 循环慢慢撸，不仅浪费时间也没效率。在一番 Google 和摸索后我找到了遍历 DataFrame 的至少 8 种方式，其中最快的和最慢的可以相差12000 倍！

本文以相加和相乘两种操作为例，测试 8 种方法的运行速度，并附上示范代码。

测试环境

  1. Macbook Pro Retina with TouchBar (13inch, 2018) i5 8GB 512GB
  2. OS: macOS Catalina 10.5.2
  3. Python 3.7.5 (default, Nov  1 2019, 02:16:23)
  4. [Clang 11.0.0 (clang-1100.0.33.8)] on darwin

示范用数据

本来想造 100 万笔的，但是 100 万笔跑 %timeit 要跑很久，最后造了 3000 笔，己经足以体现运行速度差异。为了避免快取影响，每个子实验进行前都会用这个代码重造数据。

  1. import pandas as pd
  2. import numpy as np
  4. # 生成樣例數據
  5. def gen_sample():
  6.     aaa = np.random.uniform(1,1000,3000)
  7.     bbb = np.random.uniform(1,1000,3000)
  8.     ccc = np.random.uniform(1,1000,3000)
  9.     ddd = np.random.uniform(1,1000,3000)
  10.     return pd.DataFrame({'aaa':aaa,'bbb':bbb, 'ccc': ccc, 'ddd': ddd, 'eee': None})

【Python Pandas】 遍历DataFrame的正确姿势 速度提升一万倍 - 知乎 - 图1

aaa、bbb 是本文要操作的对象；ccc、ddd 不会被操作到，只是要增大数据框的大小，模拟读入整个数据框和只读取 aaa、bbb 两列的速度差别；eee 用来存放结果

实验 1 - 两列元素相加

  1. # aaa + bbb
  2. # python 循環 + iloc 定位
  3. def method0_sum(DF):
  4.     for i in range(len(DF)):
  5.         DF.iloc[i,4] = DF.iloc[i,0] + DF.iloc[i,1]
  7. # python 循環 + iat 定位
  8. def method1_sum(DF):
  9.     for i in range(len(DF)):
  10.         DF.iat[i,4] = DF.iat[i,0] + DF.iat[i,1]
  12. # pandas.DataFrame.iterrows() 迭代器
  13. def method2_sum(DF):
  14.     for index, rows in DF.iterrows():
  15.         rows['eee'] = rows['aaa'] + rows['bbb']
  17. # pandas.DataFrame.apply 迭代
  18. def method3_sum(DF):
  19.     DF['eee'] = DF.apply(lambda x: x.aaa + x.bbb, axis=1)
  21. # pandas.DataFrame.apply 迭代 + 只讀兩列
  22. def method4_sum(DF):
  23.     DF['eee'] = DF[['aaa','bbb']].apply(lambda x: x.aaa + x.bbb, axis=1)
  25. # 列表構造
  26. def method5_sum(DF):
  27.     DF['eee'] = [ a+b for a,b in zip(DF['aaa'],DF['bbb']) ]
  29. # pandas 數組操作
  30. def method6_sum(DF):
  31.     DF['eee'] = DF['aaa'] + DF['bbb']
  33. # numpy 數組操作
  34. def method7_sum(DF):
  35.     DF['eee'] = DF['aaa'].values + DF['bbb'].values

实验 1 结果

  1. df = gen_sample()
  2. %timeit method0_sum(df)
  3. df = gen_sample()
  4. %timeit method1_sum(df)
  5. df = gen_sample()
  6. %timeit method2_sum(df)
  7. df = gen_sample()
  8. %timeit method3_sum(df)
  9. df = gen_sample()
  10. %timeit method4_sum(df)
  11. df = gen_sample()
  12. %timeit method5_sum(df)
  13. df = gen_sample()
  14. %timeit method6_sum(df)
  15. df = gen_sample()
  16. %timeit method7_sum(df)
  18. 2.06 s ± 140 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
  19. 56.9 ms ± 1.18 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
  20. 358 ms ± 24.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
  21. 93 ms ± 1.57 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
  22. 90.3 ms ± 1.16 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
  23. 1.15 ms ± 39.8 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
  24. 273 µs ± 21.6 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
  25. 130 µs ± 2.91 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

实验 2 - 两列元素相乘

  1. # aaa * bbb
  2. # python 循環 + iloc 定位
  3. def method0_times(DF):
  4.     for i in range(len(DF)):
  5.         DF.iloc[i,4] = DF.iloc[i,0] * DF.iloc[i,1]
  7. # python 循環 + iat 定位
  8. def method1_times(DF):
  9.     for i in range(len(DF)):
  10.         DF.iat[i,4] = DF.iat[i,0] * DF.iat[i,1]
  12. # pandas.DataFrame.iterrows() 迭代器
  13. def method2_times(DF):
  14.     for index, rows in DF.iterrows():
  15.         rows['eee'] = rows['aaa'] * rows['bbb']
  17. # pandas.DataFrame.apply 迭代
  18. def method3_times(DF):
  19.     DF['eee'] = DF.apply(lambda x: x.aaa * x.bbb, axis=1)
  21. # pandas.DataFrame.apply 迭代 + 只讀兩列
  22. def method4_times(DF):
  23.     DF['eee'] = DF[['aaa','bbb']].apply(lambda x: x.aaa * x.bbb, axis=1)
  25. # 列表構造
  26. def method5_times(DF):
  27.     DF['eee'] = [ a*b for a,b in zip(DF['aaa'],DF['bbb']) ]
  29. # pandas 數組操作
  30. def method6_times(DF):
  31.     DF['eee'] = DF['aaa'] * DF['bbb']
  33. # numpy 數組操作
  34. def method7_times(DF):
  35.     DF['eee'] = DF['aaa'].values * DF['bbb'].values

实验 2 结果

  1. df = gen_sample()
  2. %timeit method0_times(df)
  3. df = gen_sample()
  4. %timeit method1_times(df)
  5. df = gen_sample()
  6. %timeit method2_times(df)
  7. df = gen_sample()
  8. %timeit method3_times(df)
  9. df = gen_sample()
  10. %timeit method4_times(df)
  11. df = gen_sample()
  12. %timeit method5_times(df)
  13. df = gen_sample()
  14. %timeit method6_times(df)
  15. df = gen_sample()
  16. %timeit method7_times(df)
  18. 2.04 s ± 78.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
  19. 58.4 ms ± 3.29 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
  20. 342 ms ± 8.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
  21. 89.1 ms ± 1.59 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
  22. 90.7 ms ± 769 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
  23. 1.1 ms ± 19.9 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
  24. 263 µs ± 11 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
  25. 131 µs ± 3.35 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

速度比较

我把结果可视化，方便比较运行速度。（顺便练习画图）

【Python Pandas】 遍历DataFrame的正确姿势 速度提升一万倍 - 知乎 - 图2

可以看到最快的数组操作和最慢的 for+iloc 相比差了将近万倍，是秒级和微秒级的差别。所有方法运行速度大小关系如下（由慢至快）：for 循环 + iloc < pd.iterrows < for 循环 + at < pd.apply pd 列表构造 = np 列表构造 < pd 数组操作 < np 数组操作。pd 和 np 列表构造的速度几乎一样，np 列表构造略快一些（1.15 毫秒和 1.09 毫秒的差别），所以实验只做 pd 列表构造。

把两个秒级操作去掉，详细地比较一下

【Python Pandas】 遍历DataFrame的正确姿势 速度提升一万倍 - 知乎 - 图3

列表构造把除了数组操作以外的其他方法按在地上磨擦，数组操作把列表构造按在地上磨擦。值得注意的是，for 循环 + iat 的组合比 pandas 提供的最快遍历方法 apply 快 40% 左右，也就是说就算不懂 apply 的用法，只要把 loc/iloc 改成 at/iat，依然可以有明显的提速。另外，DataFrame 的栏位很多的时候，apply_limit 方法其实会比对对整个数据框 apply 快很多（因为不用每次读取整个数据框），只是示范数据的栏位不多所以在这里显现不出差异。

pandas 的数组操作和 numpy 的数组操作单独对比：

【Python Pandas】 遍历DataFrame的正确姿势 速度提升一万倍 - 知乎 - 图4

列表构造的运行速度是毫秒级的，数组操作是微秒级，np 数组操作是 pd 数组操作的两倍。

作图代码

  1. x = ['loop+iloc','loop+iat','iterrows','apply','apply_limit','list_build','pd_array','np_array']
  2. y_sum = [2.06*1000 , 56.9, 358, 93, 90.3, 1.15, 273/1000, 130/1000] # 單位: ms
  3. y_dot = [2.04*1000 , 58.4, 342, 89.1, 90.7, 1.10, 263/1000, 131/1000] # 單位: ms
  5. res = pd.DataFrame({'method': x*2, 'time': y_sum + y_dot, 'operation': ['sum']*len(x) + ['dot']*len(x)})
  7. import matplotlib.pyplot as plt
  8. import seaborn as sns
  9. %matplotlib inline
  10. sns.set_style('whitegrid')
  12. plt.figure(figsize=(10, 6))
  13. with sns.plotting_context("notebook", font_scale=1.5):
  14.     ax = sns.barplot(x='method', y='time', hue='operation', data=res)
  15.     plt.title('Time usage of different iterating ways')
  16.     plt.ylabel('time usage (ms)')
  18. plt.figure(figsize=(10, 6))
  19. with sns.plotting_context("notebook", font_scale=1.5):
  20.     ax = sns.barplot(x='method', y='time', hue='operation', data=res[~res.method.isin(['loop+iloc','iterrows'])])
  21.     plt.title('Time usage of different iterating ways - detailed')
  22.     plt.ylabel('time usage (ms)')
  24. res_2 = res[res.method.isin(['list_build','pd_array','np_array'])].copy()
  25. res_2['time'] = res_2['time'].values * 1000 
  27. plt.figure(figsize=(10, 6))
  28. with sns.plotting_context("notebook", font_scale=1.5):
  29.     ax = sns.barplot(x='method', y='time', hue='operation', data=res_2)
  30.     plt.title('Time usage of different iterating ways - pd_series vs np_array')
  31.     plt.ylabel('time usage (µs)')

结论

优先使用 numpy 数组操作！不能数组操作的时候用列表构造！

能用 at/iat 就不用 loc/iloc，能用 apply 就不用迭代，能用数组操作就不用其他方法。

运行速度：np 数组操作 > pd 数组操作 >>> np 列表构造 = pd 列表构造 >>> for 循环 + at > pd(分片).apply > pd.apply >>> pd.iterrows > for 循环 + iloc

关于 at/iat 和 loc/iloc 的速度比较，请见参考资料[4]。

参考资料

  1. How To Make Your Pandas Loop 71803 Times Faster
  2. rain 润：10 分钟 python 图表绘制 | seaborn 入门（二）：barplot 与 countplot
  3. seaborn.barplot
  4. 老狼：简单两步，大幅提高 python 数据处理速度
  5. seaborn.plotting_context
    https://zhuanlan.zhihu.com/p/97269320