当前位置：首页 > news >正文

076、Pandas 性能优化：从 iterrows 到 vectorize——100 倍提速的演进

news 2026/6/29 4:02:07

076、Pandas 性能优化：从 iterrows 到 vectorize——100 倍提速的演进

上周帮团队排查一个数据清洗脚本，跑了一小时还没出结果。我盯着终端里跳动的光标，心里大概有数了——八成又是哪个哥们儿在DataFrame上写了循环。打开代码一看，果然，一个iterrows套着两层if-else，处理50万行数据，每行还要做字符串拼接和条件判断。这种写法，不慢才怪。

我直接动手重写，把循环改成向量化操作，跑完只用了18秒。旁边的新人瞪大了眼睛，问我是不是换了台服务器。我说没有，只是把代码从“手动挡”换成了“自动挡”。

先看一个典型的“慢代码”长什么样

假设我们要处理一个销售订单表，根据金额和地区计算折扣后的价格。很多人会这样写：

importpandasaspdimportnumpyasnp df=pd.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,axis=1)

这段代码跑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,default=df['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

分组聚合是另一个重灾区。很多人习惯这样写：

# 慢，别这样写result=df.groupby('region').apply(lambdag:g['amount'].sum()/g['amount'].count())

直接用聚合函数：

result=df.groupby('region')['amount'].agg(['sum','count'])result['avg']=result['sum']/result['count']

或者更简洁的：

result=df.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编译：

fromnumbaimportjit@jit(nopython=True)defcumulative_returns(returns):cum=1.0result=np.empty_like(returns)fori,rinenumerate(returns):cum*=(1+r)result[i]=cumreturnresult df['cum_return']=cumulative_returns(df['daily_return'].values)