如果你的 Python 代码不够快,你有 许多编译语言选项来编写更快的扩展[1] 。在本文中,我们将重点关注 Rust,它具有以下优势:
- 现代化的工具,包括一个名为
crates.io[2] 的包仓库,以及内置的构建工具 (cargo)。 - 出色的 Python 集成和工具。用于 Python 支持的 Rust 包(它们被称为 "crates")是 PyO3[3] 。对于打包,你可以使用
setuptools-rust[4] 来集成现有的setuptools项目,或者对于独立扩展,你可以使用 Maturin[5] 。 - 内存和线程安全,因此与 C 和 C++ 相比,更不容易崩溃或内存损坏。
具体来说,我们将:
- 用 Python 实现一个小算法。
- 将其重新实现为 Rust 扩展。
- 优化 Rust 版本,使其运行更快。
近似计算唯一值的数量
作为一个激励性的例子,我们将研究计算列表中有多少个唯一值的任务。如果我们想得到一个精确的答案,这很容易实现:
def count_exact(items): return len(set(items))
set() 只能包含一个项目一次,所以这会去重所有项目,然后计算它们。
这个解决方案的问题是内存使用。如果我们有 10,000,000 个唯一值,我们将创建一个包含 10,000,000 个条目的集合,这将使用相当多的内存。
因此,如果我们担心内存使用,我们可以使用一种概率算法,给我们一个_近似_答案。在许多情况下,这就足够了,而且它可以使用少得多的内存。我们将使用 Chakraborty、Vinodchandran 和 Meel 的 一个非常简单的算法[6] :
import random def count_approx_python(items): s = set() for item in items: if random.random() < 0.1: s.add(item) return len(s) * 10
运行一个例子
让我们看一组示例单词:
import random words = [f"word{i}" for i in range(100_000)] random.shuffle(words)
我们有 100,000 个不同的单词,这意味着 count_exact() 将创建一个大小为 100,000 的集合。同时,count_approx() 将有一个大小为 1739 的集合。它最终偶尔会同时有两个集合,但它仍然只会使用准确算法 3% 的内存。
我们可以比较结果:
print("Exact:", count_exact(words)) print("Approximate:", count_approx_python(words))
我连续运行了三次,得到:
Exact: 100000 Approximate: 99710 Exact: 100000 Approximate: 99880 Exact: 100000 Approximate: 100090
近似版本有所变化,但它非常接近 - 同时使用的内存要少得多。
速度比较
我们可以比较两种实现的速度:
import time start = time.time() count_exact(words) print("Exact:", time.time() - start) start = time.time() count_approx_python(words) print("Approximate:", time.time() - start)
这给我们:
Exact: 0.015625 Approximate: 0.078125
我们的新函数使用更少的内存,但速度慢了 5 倍。 让我们看看是否可以通过用 Rust 重写 count_approx_python() 来加速它。
加速:创建我们的 Rust 项目
使用 Maturin[4] Python 打包工具,我们可以非常快速地创建一个新的 Rust Python 扩展,并使用 PyO3[2] 可以轻松地与 Python 对象交互。
1. 使用 Maturin 初始化项目
你可以 pipx 或 pip install maturin,然后你可以使用 Maturin 来初始化一个全新的基于 PyO3 的项目:
$ maturin init 🤷 What kind of project do you want to create? pyo3 🤷 What's the name of your Python package? approx_count [...] ✨ Done! New project created approx_count
这创建了我们需要的所有文件,用于一个基本的基于 Rust 的 Python 包:
approx_count/
├── Cargo.toml
├── pyproject.toml
└── src
└── lib.rs
我们可以从一开始就 pip install 这个包,不需要进一步的设置:
$ pip install .
2. 添加依赖项
Rust 不包含内置的随机数生成库,所以我们将使用 Rust 的构建/包管理器 cargo 添加一个第三方 crate(Rust 术语中的可安装包):
$ cargo add rand Updating crates.io index Adding rand v0.8.5 to dependencies. Features: + std + std_rng + alloc - custom_derive - getrandom - libc - log - min_const_gen - packed_simd - serde - serde1 - simd_support Updating crates.io index
这更新了 Cargo.toml 文件,其中列出了构建我们代码所需的 Rust 依赖项。相关部分现在看起来像这样:
[dependencies] pyo3 = "0.19.0" rand = "0.8.5"
当我们初始化项目模板时,Maturin 自动添加了 pyo3 依赖项。
命令输出中的功能列表是可以在编译时启用以添加更多功能的标志;我们稍后会使用其中一个。
3. 第一个 Rust 版本
现在,我们需要更新 src/lib.rs 中的 Rust 代码来实现我们的函数:
use pyo3::prelude::*; use rand::random; #[pyfunction] fn count_approx_rust(items: Vec<PyObject>, py: Python<'_>) -> PyResult<usize> { let mut s = pyo3::types::PySet::empty(py)?; for item in items { if random::<f64>() < 0.1 { s.add(item)?; } } Ok(s.len() * 10) } #[pymodule] fn approx_count(_py: Python, m: &PyModule) -> PyResult<()> { m.add_function(wrap_pyfunction!(count_approx_rust, m)?)?; Ok(()) }
4. 测量性能
我们的新版本在速度方面如何比较?
首先,我们安装我们的新包:
现在我们可以从 Python 代码中导入我们的新函数并测量其速度:
from approx_count import count_approx_rust import time start = time.time() count_approx_rust(words) print("Rust (naive):", time.time() - start)
以下是新版本的比较:
| 版本 | 耗时(秒) |
|---|---|
| Python | 0.78 |
| Rust (naive) | 0.37 |
它的速度是两倍。为什么不更快?
这在逻辑上与 Python 实现完全相同的代码,只是用 Rust 实现,而且稍微更冗长一些。我们仍然以相同的方式与 Python 对象交互,迭代 Python 列表,并广泛与 Python 集合交互。所以代码的那部分不会有任何不同的运行方式。
让我们更快,第 2 部分:优化 Rust 代码
接下来我们将以四种不同的方式优化我们的代码,我分别测量了所有这些方式都提高了性能。
优化 1:链接时优化
首先,我们将启用链接时优化,Rust 编译器在编译过程的后期(在链接期间)优化代码。这意味着编译速度较慢,但通常会产生运行速度更快的代码。我们通过在 Cargo.toml 中添加以下内容来实现这一点:
[profile.release] lto = true
优化 2 和 3:更快的随机数生成
我们还从使用 rand::random() (它使用线程本地随机数生成器(RNG))切换到我们自己管理的 RNG,消除了线程本地查找的开销。
同时,我们切换到使用比默认 RNG 更快的 "small" RNG;它在防御哈希拒绝服务攻击方面不太安全,但对我们的目的来说可能没问题。为此,我们在 Cargo.toml 中为 rand crate 添加 smallrng 功能:
rand = { version = "0.8.5", features = ["small_rng"] }
我们需要修改代码,我们将在下面看到。
优化 4:仅存储哈希值
最后,我们从在基于 Python 的集合中存储 Python 对象切换到仅存储 Python 对象的哈希值,该哈希值由 obj.__hash__() 计算。如果两个对象的哈希值相同会发生什么?结果将偏差 1。
我们已经在使用概率函数,所以我们已经可以接受略微错误的结果。冲突应该很少见,如果我们有许多唯一值,偶尔出现一个偏差并不重要;99712 并不比 99713 更错。
由于我们不再存储 Python 对象,我们可以切换到使用 Rust 的 std::collections::HashSet,它有一些更好的 API,可能会稍微快一些。然而,Rust HashSet 会想要对值进行哈希...而我们不想再次对它们进行哈希,它们已经是预先哈希过的。
为了避免双重哈希,我们还添加了 Rust crate nohash_hasher:
$ cargo add nohash-hasher
这将允许我们使用基于 Rust 的 HashSet,它假设它存储的值已经是哈希值,可以按原样在 HashSet 中使用,无需进一步哈希。Cargo.toml 依赖项部分现在看起来像这样:
[dependencies] pyo3 = "0.19.0" rand = { version = "0.8.5", features = ["small_rng"] } nohash-hasher = "0.2.0"
我们还需要更新我们的代码来使用它,我们接下来会看到。
我们的新代码,包含所有 4 个优化
这是我们更新后的代码:
use pyo3::prelude::*; use rand::rngs::SmallRng; use rand::{Rng, SeedableRng}; use std::collections::HashSet; use nohash_hasher::IntSet; #[pyfunction] fn count_approx_rust(items: Vec<PyObject>, py: Python<'_>) -> PyResult<usize> { let mut rng = SmallRng::from_entropy(); let mut s: IntSet<u64> = IntSet::default(); for item in items { if rng.gen::<f64>() < 0.1 { let hash = item.as_ref(py).hash()?; s.insert(hash as u64); } } Ok(s.len() * 10) } #[pymodule] fn approx_count(_py: Python, m: &PyModule) -> PyResult<()> { m.add_function(wrap_pyfunction!(count_approx_rust, m)?)?; Ok(()) }
同样,我们可以通过以下方式安装更新后的版本:
以下是我们优化版本与之前版本相比运行所需的时间:
| 版本 | 耗时(秒) |
|---|---|
| Python | 0.78 |
| Rust (naive) | 0.37 |
| Rust (optimized) | 0.21 |
为什么这不更快,以及其他想法
为什么 Rust 代码不更快?我们最新的优化版本仍然与 Python 列表交互,即 items,并使用 Python 的 __hash__ API 来对每个对象进行哈希。这意味着我们仍然受限于这两个交互的 Python API 的速度。如果我们将项目作为 Arrow 列或整数的 NumPy 数组传入,我们可能会运行得更快。
更广泛地说,如果我们想要更快,我们还应该考虑使用不同的算法。有许多近似计数算法,我选择这个只是因为它简单易懂,而不是因为它特别快。
大局观:为什么选择 Rust?
Rust 允许我们通过给我们访问编译语言的能力来加速我们的代码。但这对 C 或 C++ 或 Cython 也是如此。
然而,与这些语言不同,Rust 有现代化的工具,内置包管理器和构建系统:
- 添加新的依赖项就像
cargo add <cratename>一样简单。浏览可用 crate 的好网站是lib.rs[7] 。 - 不太明显,但仍然重要:这段代码将在 macOS 和 Windows 上编译,无需额外工作。相比之下,跨平台管理 C 或 C++ 依赖项可能非常痛苦。
Rust 还有出色的 Python 集成,可以高级访问 Python API,并且有易于使用的打包工具,如 Maturin。
Rust 既可以扩展到小项目 - 如本文中所示 - 也可以扩展到更大的项目,这要归功于内存和线程安全性以及其强大的类型系统。Polars 项目有一个 330K 行的通用 Rust 核心,然后用更多的 Rust 和 Python 包装,制成 Python 扩展。
下次你考虑加速一些 Python 代码时,试试 Rust。它不是一种简单的语言,但值得你花时间学习。
参考链接
- 许多编译语言选项来编写更快的扩展: https://pythonspeed.com/articles/rust-cython-python-extensions/
crates.io: https://crates.io/- PyO3: https://pyo3.rs/
setuptools-rust: https://setuptools-rust.readthedocs.io/- Maturin: https://www.maturin.rs/
- 一个非常简单的算法: https://arxiv.org/abs/2301.10191
lib.rs: https://lib.rs/
