076、Pandas 性能优化:从 iterrows 到 vectorize——100 倍提速的演进
076、Pandas 性能优化从 iterrows 到 vectorize——100 倍提速的演进上周帮团队排查一个数据清洗脚本跑了一小时还没出结果。我盯着终端里跳动的光标心里大概有数了——八成又是哪个哥们儿在DataFrame上写了循环。打开代码一看果然一个iterrows套着两层if-else处理50万行数据每行还要做字符串拼接和条件判断。这种写法不慢才怪。我直接动手重写把循环改成向量化操作跑完只用了18秒。旁边的新人瞪大了眼睛问我是不是换了台服务器。我说没有只是把代码从“手动挡”换成了“自动挡”。先看一个典型的“慢代码”长什么样假设我们要处理一个销售订单表根据金额和地区计算折扣后的价格。很多人会这样写importpandasaspdimportnumpyasnp dfpd.DataFrame({amount:np.random.uniform(100,10000,500000),region:np.random.choice([华东,华北,华南,西南],500000)})defcalculate_discount(row):ifrow[region]华东:ifrow[amount]5000:returnrow[amount]*0.85else:returnrow[amount]*0.9elifrow[region]华北:returnrow[amount]*0.88else:returnrow[amount]*0.95# 别这样写慢到怀疑人生df[discounted]df.apply(calculate_discount,axis1)这段代码跑50万行在我的机器上大概需要12秒。看起来还行别急如果换成iterrows直接奔着40秒去了。更可怕的是如果业务逻辑再复杂一点比如嵌套几个字典查找、正则匹配半小时都跑不完。为什么循环和apply这么慢这里踩过坑的人都知道iterrows返回的是Series对象每次迭代都要做类型推断、索引对齐Python解释器在每一行都要重新进入Pandas的C扩展层。apply虽然看起来高级一点本质上还是在Python层面逐行调用函数没有利用到NumPy底层的向量化能力。打个比方循环就像你一个一个地搬砖向量化操作就像用铲车一次性铲起一堆砖。CPU的SIMD指令集就是那把铲车但前提是你得把数据组织成它能理解的形式。第一步用向量化操作替换条件逻辑对于上面的折扣计算最直接的做法是用np.select或者布尔索引# 这才是正确姿势conditions[(df[region]华东)(df[amount]5000),(df[region]华东)(df[amount]5000),(df[region]华北)]choices[df[amount]*0.85,df[amount]*0.9,df[amount]*0.88]# 默认值给0.95df[discounted]np.select(conditions,choices,defaultdf[amount]*0.95)这段代码跑完只需要0.3秒。40倍提速而且代码更短、更清晰。np.select会一次性生成所有条件的布尔掩码然后通过C级别的循环完成赋值完全没有Python层面的逐行开销。第二步字符串操作也要向量化很多人处理字符串时习惯用apply加lambda# 慢别这样写df[clean_region]df[region].apply(lambdax:x.replace(华,中))换成Pandas自带的字符串方法df[clean_region]df[region].str.replace(华,中)str访问器背后调用的是NumPy的向量化字符串操作速度能快10倍以上。如果要做正则匹配用str.contains、str.extract不要自己写循环。第三步groupby之后别用apply分组聚合是另一个重灾区。很多人习惯这样写# 慢别这样写resultdf.groupby(region).apply(lambdag:g[amount].sum()/g[amount].count())直接用聚合函数resultdf.groupby(region)[amount].agg([sum,count])result[avg]result[sum]/result[count]或者更简洁的resultdf.groupby(region)[amount].mean()groupby的聚合操作是高度优化的C代码而apply会把每个分组的数据传到Python层来回切换上下文性能损失巨大。第四步终极武器——用NumPy的ufunc如果Pandas没有提供你需要的向量化函数别急着写循环。看看能不能用NumPy的通用函数ufunc组合实现。比如我们要对金额做分段标记小于1000为“低”1000-5000为“中”大于5000为“高”。# 用np.selectbins[0,1000,5000,np.inf]labels[低,中,高]df[level]np.select([df[amount]1000,(df[amount]1000)(df[amount]5000),df[amount]5000],labels)或者用pd.cut但注意pd.cut内部也是向量化的df[level]pd.cut(df[amount],bins[0,1000,5000,np.inf],labels[低,中,高])什么时候真的需要用循环说了这么多向量化的好处但有些场景确实绕不开循环。比如每一行的计算依赖上一行的结果比如累计收益计算需要调用外部API或数据库网络IO无法向量化复杂的业务规则无法用简单的数学表达式描述对于第一种情况可以用numba加速。Pandas 0.24之后支持pd.Series.rolling配合apply但更推荐用numba的JIT编译fromnumbaimportjitjit(nopythonTrue)defcumulative_returns(returns):cum1.0resultnp.empty_like(returns)fori,rinenumerate(returns):cum*(1r)result[i]cumreturnresult df[cum_return]cumulative_returns(df[daily_return].values)numba会把Python循环编译成机器码速度接近C语言。但注意numba只支持NumPy数组和基本Python类型不能直接传DataFrame。一个真实的优化案例上个月处理一个用户行为日志200万行需要根据时间戳和用户ID计算会话间隔。原始代码用iterrows嵌套groupby跑了25分钟。我重构后用了三步先用sort_values按用户和时间排序用groupby加shift获取上一行时间戳向量化计算时间差最终代码跑了2.3秒。那个同事后来跟我说他以为Pandas就只能那么慢。个人经验总结别把Pandas当成Excel的Python版本来用。Pandas的底层是NumPy和C扩展它的设计哲学就是“批量操作”。你每写一个for row in df.iterrows()就是在强迫Pandas放弃它的优势退化成纯Python的速度。调试性能问题时我的习惯是先跑一个df.info()看数据类型确保没有object类型的数值列。然后用%timeit测试关键操作如果单次操作超过1秒就考虑向量化。如果向量化实在做不到再考虑numba或者cython。最后说一句不要为了炫技而用向量化。如果数据量只有几千行循环完全够用没必要为了追求极致性能把代码写成天书。但如果你在处理几十万行以上的数据向量化不是可选项是必选项。