让我们构建和优化一个用于 Python 的 Rust 扩展

主要观点:当 Python 代码不够快时,可选择编译语言编写更快的扩展,本文重点介绍 Rust。Rust 有现代工具(如crates.iocargo)、优秀的 Python 集成工具(如PyO3setuptools-rustMaturin),且内存和线程安全。文中以计数唯一值为例,先在 Python 中实现,再重写为 Rust 扩展并优化,最终使 Rust 版本运行更快。具体如下:

  • 计数唯一值(近似方式):Python 中精确计数用set,但内存占用大,可使用概率算法count_approx_python,示例中 100,000 个不同单词,精确算法需 100,000 内存,近似算法只需 1739 内存,且结果接近但速度慢(Python 0.78 秒,Rust 朴素版 0.37 秒)。
  • 创建 Rust 项目:使用Maturin初始化项目,创建基本文件,pip install即可安装,添加rand依赖,更新src/lib.rs中的 Rust 代码实现函数count_approx_rust,测量性能发现比 Python 快两倍。
  • 优化 Rust 代码:通过四种方式优化,包括启用链接时优化(在Cargo.toml中添加[profile.release] lto = true)、使用更快的随机数生成(从rand::random()切换到自己管理的SmallRng,并添加small_rng特性)、只存储 Python 对象的哈希值(使用nohash_hasher避免重复哈希),优化后的 Rust 代码运行时间缩短为 0.21 秒,但仍受限于与 Python 列表的交互速度,若传入其他数据结构可能更快,且可考虑使用其他更快的近似计数算法。
  • 整体优势:Rust 能通过编译语言提速,有现代工具(如cargo add方便添加依赖,lib.rs可浏览可用 crate),跨平台编译方便,Python 集成好,可用于小项目和大项目,如 Polars 项目。下次提速 Python 代码可尝试 Rust。
阅读 28
0 条评论