首页
学习
活动
专区
圈层
工具
发布

7 个 pandas 矢量化写法实操:算术运算、条件判断、分组聚合怎么写才快

如果你了解 pandas 那么大概都听过那句被反复念叨的规则:"永远不要逐行遍历 DataFrame"。但这代价到底有多大,很少有人说清楚——.apply()、np.where()、.map()、df.eval()各管一摊,节省的时间也差得远,不能随手挑一个就上。

与其再写一篇"矢量化更快"式的文章,不如实测一遍:构建一个 20 万行、电商风格的数据集,把七个常见任务的循环版本和矢量化版本分别计时。得到的结果有些结果和预期一致,有一个结果确实出乎意料。

下面是数据、代码,以及七个值得记住的技巧。

pandas 循环慢在哪

pandas 的 Series 本质上是对 NumPy 数组的封装。用for i in range(len(df))或.iterrows()逐行遍历时,就脱离了 NumPy 编译好的 C 语言循环,退回到 Python 解释器——每一行都要经历类型检查、对象创建、函数调用开销,行数一多这些开销就会累积起来。

矢量化操作跳过了这些,把整个运算下推到 C 语言层面,一次性处理整个数组。下文所有内容都只是"让 pandas/NumPy 用 C 去循环,而不是自己用 Python 去循环"的不同变体。

本文全篇使用的基准测试数据集:

import numpy as np

import pandas as pd

np.random.seed(42)

N = 200_000df = pd.DataFrame({

  "price": np.random.uniform(10, 500, N).round(2),

  "quantity": np.random.randint(1, 20, N),

  "category": np.random.choice(

      ["electronics", "grocery", "apparel", "toys"], N

  ),

  "customer_id": np.random.randint(1, 5000, N),

  "raw_name": [f"  Product_{i}  " for i in np.random.randint(1, 1000, N)],

})

下文每一个计时结果,都是在该数据集上多次运行取得的中位数。

1、矢量化运算代替逐行数学计算

循环写法:

totals = []

for i in range(len(df)):

  totals.append(df["price"].iloc[i] * df["quantity"].iloc[i])

替代写法:

df["total"] = df["price"] * df["quantity"]

这是最常见方法,同时差距也最悬殊:循环耗时 9.84 秒,矢量化写法只需 0.0014 秒,快了约 7075 倍。循环里的.iloc[i]尤其费性能,每次调用都要重新校验索引、重新封装返回结果,这还是叠加在每行本身的解释器开销之上的。如果发现自己在for循环里对两列做算术,这是最简单的提速方式

2、用 np.where() 代替 if/elif 链

循环写法:

tiers = []

for p in df["price"]:

  if p > 300:

      tiers.append("premium")

  elif p > 100:

      tiers.append("standard")

  else:

      tiers.append("budget")

替代写法:

df["tier"] = np.where(

  df["price"] > 300, "premium",

  np.where(df["price"] > 100, "standard", "budget")

)

嵌套的np.where()把条件计算成布尔数组逐元素选取对应的值,全程在 C 语言层面完成。测得结果:矢量化写法 0.0075 秒,循环写法 0.022 秒,提速 3 倍。这个幅度比第 1 条小得多,所以一旦涉及字符串,np.where的条件分支运算就拿不到那种爆炸式收益,因为比较操作本身在对象数组上会带来额外开销。分支超过两三个的话,np.select()更简洁,可读性也更好。

3、用 .map() 代替循环中的字典查找

循环写法:

discount_lookup = {

  "electronics": 0.10, "grocery": 0.02,

  "apparel": 0.15, "toys": 0.05,

}

discounts = [discount_lookup[c] for c in df["category"]]

替代写法:

df["discount"] = df["category"].map(discount_lookup)

.map()正是为这种场景设计的:把一列数据通过字典、Series 或函数转换。测得结果是矢量化写法 0.0084 秒,列表推导式 0.065 秒,快了 7.7 倍;类别数量超过两三个时,可读性也比循环或 if/elif 阶梯好。同样的方法也能用来重命名取值、重新编码问卷答案,或者在不做完整merge()的情况下拼接一张小的查找表。

4、用 .str 访问器代替字符串遍历循环

循环写法:

clean_names = [name.strip().lower() for name in df["raw_name"]]

替代写法:

df["clean_name"] = df["raw_name"].str.strip().str.lower()

.str访问器测得结果:矢量化写法 0.076 秒,循环写法 0.089 秒,只快了 1.2 倍。这个结果可能出乎意料,因为毕竟本以为字符串操作也能像数值操作那样拿到量级上的加速。事实不完全如此:pandas 的.str方法底层仍然是逐元素地在 Python 层面迭代,因为 NumPy/pandas 把字符串存成 Python 对象,不是固定宽度的 C 语言类型。代码会更干净、更方便链式调用,也自带 NaN 处理,但拿不到数值列那种数量级的提速。如果字符串处理真是大数据集上的瓶颈,可以先把该列转成 pandas 的category类型(取值大量重复时适用),或者用较新的 PyArrow 支持的字符串类型(pd.array(..., dtype="string[pyarrow]")),会更接近真正的矢量化性能。

5、用 groupby().transform() 代替遍历分组

循环写法:

cat_totals = df.groupby("category")["total"].sum().to_dict()

shares = [

  tot / cat_totals[cat]

  for cat, tot in zip(df["category"], df["total"])

]

替代写法:

df["share_of_category"] = (

  df["total"] / df.groupby("category")["total"].transform("sum")

)

.transform()一步完成分组聚合值的计算并广播回原始行形状,不需要手动建字典,也不需要手动 zip。测得结果是矢量化写法 0.019 秒,循环写法 0.119 秒,快了约 6.2 倍。分组数量越多,它比手写循环的伸缩性越好,因为 pandas 只需为每个分组计算一次聚合值。

6、用 pd.cut() 代替手写分箱逻辑

循环写法:

bands = []

for p in df["price"]:

  if p <= 50:

      bands.append("low")

  elif p <= 150:

      bands.append("mid")

  elif p <= 300:

      bands.append("high")

  else:

      bands.append("luxury")

替代写法:

bins = [0, 50, 150, 300, 500]

labels = ["low", "mid", "high", "luxury"]

df["price_band"] = pd.cut(df["price"], bins=bins, labels=labels)

pd.cut()(以及按频率分箱的pd.qcut())彻底取代了手写分桶逻辑,返回的是内存效率更高的Categorical列,而不是字符串列表。测得结果是矢量化写法 0.0066 秒,循环写法 0.029 秒,快了 4.4 倍;真正的收益体现在可维护性上——改分箱边界只是一行代码的事,不用重写整个 if/elif 链。

7、用 df.query() + .eval() 代替循环中的过滤与计算

循环写法:

result = [

  p * q * 0.9

  for p, q in zip(df["price"], df["quantity"])

  if p > 200 and q > 5

]

替代写法:

result = (

  df.query("price > 200 and quantity > 5")

    .eval("price * quantity * 0.9")

)

query()和eval()解析字符串表达式,交给 pandas 自身的表达式引擎运行(可选由numexpr提供底层支持),避免创建多个中间布尔数组和临时列。测得结果是矢量化写法 0.018 秒,循环写法 0.040 秒,快了约 2.2 倍。更大的好处体现在多条件过滤时的可读性上:df.query("price > 200 and quantity > 5 and category == 'electronics'")读起来更接近日常英语,比起用&、|运算符、每个子句都要加括号的写法清楚得多。

七种技巧对比

数值运算这一项完全在另一个量级——不管是不是对数坐标,7075 倍的提速都让其他结果相形见绌。中间那一档(.map()、groupby().transform()、pd.cut())带来的是稳定的数倍提速。.str访问器则收益不明显:虽然净收益依然为正,但提醒一点——一旦涉及字符串对象,"矢量化"不等于"快"。

速查表

总结

纯数值运算是 pandas 展现速度优势最明显的地方;条件逻辑、查找和分组操作带来的是稳定、有意义的提升;字符串操作则是唯一需要放低预期的地方——文本处理真成了瓶颈时,category类型或 PyArrow 支持的字符串类型是现成的解法。

下次敲下for i in range(len(df)):时,不妨先停一下,想想问题到底属于这七种形状中的哪一种,而不是条件反射式地掏出.apply()草草了事。

作者:Himani

点个在看你最好看!

  • 发表于:
  • 原文链接https://page.om.qq.com/page/On3I1KC0qIoQd59avQu58vvA0
  • 腾讯「腾讯云开发者社区」是腾讯内容开放平台帐号(企鹅号)传播渠道之一,根据《腾讯内容开放平台服务协议》转载发布内容。
  • 如有侵权,请联系 cloudcommunity@tencent.com 删除。
领券