如何使用 Rust 和 PyO3 加速 Python

Python 以其简洁性和丰富的第三方包生态而著称，但涉及重型计算、解析 GB 级数据或验证数百万个对象时，解释器就会开始变慢。以前遇到这种情况，开发者必须手动编写 C API 扩展。这个过程非常耗费人力：调试内存泄漏和手动引用管理很快就会消磨掉开发的乐趣。

几年前出现了一条更安全的道路。PyO3 库支持用 Rust 编写原生模块，同时保持严格的内存安全保证和清晰的抽象。

如果你使用过 Pydantic v2 验证器、orjson JSON 解析器、Polars 框架或 OpenAI 的 Tiktoken 分词器，你就已经运行过用 PyO3 创建的代码了。

为什么要将 Rust 与 Python 结合

Python 擅长业务逻辑和快速原型开发。Rust 在需要最高速度、低资源消耗和并发性的场景中胜出。

PyO3 解决了两个问题：

1. 将 Python 应用程序的性能瓶颈卸载到独立的 Rust 函数或结构中。
2. 在 Rust 二进制文件中嵌入 Python 解释器以执行脚本。

大多数情况下，该项目用于第一种场景。

如何在实践中编写模块

对于将 Rust 代码构建为 Python 包，作者推荐 maturin 工具。它处理编译并创建 wheel 文件。

让我们创建一个简单的模块，接收两个数字相加并以字符串形式返回结果。

首先，设置虚拟环境：

mkdir string_sum && cd string_sum
python -m venv .env
source .env/bin/activate
pip install maturin
maturin init --bindings pyo3

在 src/lib.rs 文件中的 Rust 端逻辑：

use pyo3::prelude::*;

#[pyfunction]
fn sum_as_string(a: usize, b: usize) -> PyResult<String> {
    Ok((a + b).to_string())
}

#[pymodule]
mod string_sum {
    use super::*;

    #[pymodule_export]
    fn init(m: &Bound<'_, PyModule>) -> PyResult<()> {
        m.add_function(wrap_pyfunction!(sum_as_string, m)?)?;
        Ok(())
    }
}

通过 #[pyfunction] 和 #[pymodule] 宏自动生成 Python C API 代码。开发者无需手动将 PyObject 类型转换为 Rust 结构体，也无需反向转换。

开始构建：

maturin develop

现在可以直接从 Python 调用该函数：

import string_sum

result = string_sum.sum_as_string(5, 20)
print(result)  # '25'

如果想在生产环境中测量真实性能，只需使用优化标志构建模块：maturin develop --release。

实际应用案例

在 Pydantic V2 中，所有类型检查和验证逻辑都在单独的包 pydantic-core 中用 Rust 重写了。性能得到了显著提升。

Polars 库凭借 Rust 中的并行计算，正在大型数据集处理中取代 Pandas。用户编写熟悉的 Python 代码，而重型工作则由编译后的模块完成。

在机器学习中，Hugging Face 在 tokenizers 包中使用了 PyO3，OpenAI 在 tiktoken 中也使用了它。在发送到神经网络之前进行快速文本解析显著降低了服务延迟。

配套工具

围绕 PyO3 已经形成了一套实用的工具集：

• maturin 无需编写复杂的 setup.py 配置。
• rust-numpy 将 NumPy 数组与 Rust 绑定，无需在内存中复制数据。
• pyo3-async-runtimes 将 asyncio 与 Tokio 异步运行时连接起来。
• pythonize 使用 Serde 转换数据结构。

结论

如果你的项目存在性能瓶颈，PyO3 可以帮助你解决问题，无需切换到 C 或 C++。这里的主要困难在于学习 Rust 本身。需要理解所有权概念和数据类型，但该库完全处理了与 CPython 交互的常规工作。