我的个人技术笔记

21.2 加速工具

4 min read

核心内容

借助 NumPy 向量化、bottleneck 加速、numba JIT 编译 与 PyArrow 字符串 四类工具大幅提升 pandas 计算性能。

一、NumPy 加速

pandas 底层基于 NumPy，直接使用 NumPy 函数可避免 pandas 层开销。

import pandas as pd
import numpy as np
 
df = pd.DataFrame(np.random.rand(1_000_000, 5), columns=list('abcde'))
 
# pandas 方法
result = df['a'] * 2 + np.sin(df['b'])
 
# NumPy 直接操作（当无缺失值时可更高效）
arr = df[['a', 'b']].to_numpy()
result_np = arr[:, 0] * 2 + np.sin(arr[:, 1])
 
# 使用 NumPy ufunc
df['a'].values.mean()          # 比 df['a'].mean() 略快
np.log1p(df['a'].values)       # 计算 log(1+x)

适用场景

  • 纯数值、无缺失值的大数组运算
  • 频繁调用的自定义数值函数
  • 矩阵运算：df.to_numpy() @ other

二、bottleneck 加速

bottleneck 是 C 语言实现的 NumPy 加速库，pandas 默认开启。

# 确认是否开启
pd.get_option('compute.use_bottleneck')   # True
 
# 手动开关
pd.set_option('compute.use_bottleneck', True)
pd.set_option('compute.use_bottleneck', False)

加速的运算：sum、mean、std、var、min、max、median、rank、rolling 相关聚合等。

安装 

pip install bottleneck

三、numba 加速

numba 通过 LLVM 将 Python 函数 JIT 编译为机器码，适合无法向量化的循环逻辑。

import pandas as pd
import numpy as np
from numba import jit
 
# 普通 Python 函数
def compute_python(x, y):
    total = 0.0
    for i in range(len(x)):
        total += x[i] * y[i]
    return total
 
# numba JIT 编译
@jit(nopython=True)
def compute_numba(x, y):
    total = 0.0
    for i in range(len(x)):
        total += x[i] * y[i]
    return total
 
x = np.random.rand(1_000_000)
y = np.random.rand(1_000_000)
compute_numba(x, y)      # 首次调用编译，后续极快

与 pandas 结合

# 1. rolling.apply 使用 numba 引擎
df['b'].rolling(10).apply(
    lambda x: np.argmax(x), engine='numba', raw=True
)
 
# 2. groupby.transform / apply 使用 numba
df.groupby('g')['v'].transform(
    lambda x: x - x.mean(), engine='numba'
)
 
# 3. 自定义函数中调用 @jit 函数
@jit(nopython=True)
def rolling_custom(arr):
    return arr.mean() + arr.std()
 
df['c'] = df['b'].rolling(10).apply(rolling_custom, raw=True)

numba 限制

  • 需要安装 numba：pip install numba
  • 首次调用有编译开销（可在小数据上预热）
  • nopython=True 模式下不支持 pandas 对象，需传 NumPy 数组

四、PyArrow 字符串性能

PyArrow 字符串后端在内存占用与运算速度上均优于 object 字符串。

import pandas as pd
 
# 方式 1：全局开启
pd.set_option('future.infer_string', True)   # 2.0+ / 2.3 中推荐
 
# 方式 2：显式 dtype
df['name'] = df['name'].astype('string[pyarrow]')
 
# 方式 3：读取时指定
df = pd.read_csv('data.csv', dtype_backend='pyarrow')
 
# 方式 4：ArrowDtype
import pyarrow as pa
df['name'] = df['name'].astype(pd.ArrowDtype(pa.string()))

性能对比示例

df = pd.DataFrame({
    'name': ['apple', 'banana', 'cherry'] * 300_000,
    'value': range(900_000)
})
 
# object 字符串
df_obg = df.astype({'name': 'object'})
%timeit df_obg['name'].str.contains('a')      # 慢
 
# pyarrow 字符串
df_pa = df.astype({'name': 'string[pyarrow]'})
%timeit df_pa['name'].str.contains('a')       # 快 2-5 倍
 
# 内存对比
df_obg.memory_usage(deep=True)
df_pa.memory_usage(deep=True)

PyArrow 字符串支持的更多操作

df['name'].str.upper()
df['name'].str.startswith('a')
df['name'].str.contains('a', regex=True)
df['name'].str.extract(r'(\w+)')

与 pandas 字符串访问器差异

PyArrow 字符串的缺失值统一为 pd.NA，行为更一致；str 访问器绝大部分方法可用，个别正则高级特性可能受限。

加速策略选择

场景推荐工具
数值向量化运算NumPy ufunc / 内置方法
滚动、聚合大数组bottleneck
复杂循环 / 自定义算法numba（@jit + nopython=True）
大量字符串处理PyArrow 字符串 dtype
表达式过滤eval() / query() + numexpr

相关笔记

Graph View

Backlinks