Rust 使 Python 函数速度提高 5000%

大家应该都听说过，Rust 因其卓越的性能和安全性，正被越来越多的科技巨头采用，推荐开发者使用Rust来构建关键软件。

今天，来深入学习一下，如何利用 Rust 来大幅提升你的 Python代码性能！

寻找第N个质数，Python实现

下列，编写一个 Python 函数来寻找第N个质数。

def is_prime(num):
    # 检查一个数字是否为质数。
    if num < 2:
        return False
    for i in range(2, num):
        if num % i == 0:
            return False
    return True

def find_nth_prime(n):
    # 寻找第N个质数。
    count = 0
    num = 1
    while count < n:
        num += 1
        if is_prime(num):
            count += 1
    return num
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

python prime_finder.py 50 当我在终端中执行命令时，获得的响应如下所示。

寻找第 N 个素数，Rust 实现

fn is_prime_rust(num: u32) -> bool {
    if num < 2 {
        return false;
    }
    for i in 2..num {
        if num % i == 0 {
            return false;
        }
    }
    true
}

fn find_nth_prime_rust(n: u32) -> u32 {
    let mut count: u32 = 0;
    let mut num: u32 = 1;

    while count < n {
        num += 1;

        if is_prime_rust(num) {
            count += 1;
        }
    }

    num
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27

在 Python 中集成 Rust 代码

步骤1：初始化我们的项目

mkdir rust_prime_funcs && cd rust_prime_funcs # 创建项目目录

python -m venv env # 创建Python虚拟环境

source ./env/Scripts/activate # 激活虚拟环境
1
2
3
4
5

第二步：安装Maturin

接下来，我们使用Maturin，这是一个帮助我们构建和发布带有pyo3、rust-cpython和CFFI绑定以及Rust二进制文件作为Python包的工具。

pip install maturin  # 使用pip安装maturin
1

步骤3：使用Rust创建Python模块

在这一步中，我们使用Rust函数创建一个Python模块。

// rust_prime_funcs/src/lib.rs

use pyo3::prelude::*;

// Rust中的辅助函数
#[pyfunction]
fn is_prime_rust(num: u32) -> bool {
    // ...（与上面Rust实现相同）
}

// Rust中的第N个质数查找函数
#[pyfunction]
fn find_nth_prime_rust(n: u32) -> u32 {
    // ...（与上面Rust实现相同）
}

/// 用Rust实现的Python模块。
#[pymodule]
fn rust_prime_funcs(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(find_nth_prime_rust, m)?)?;
    Ok(())
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

第四步：构建Python模块

我们使用以下命令在终端构建Python模块。

maturin develop --release
1

步骤5：比较Python与Rust函数的性能

最后，我们比较两个函数的性能

# rust_prime_funcs/test/prime_finder.py

import sys
from timeit import timeit

from rust_prime_funcs import find_nth_prime_rust

# ...（Python代码与上面相同）

def main():
    n = int(sys.argv[1])

    # ...（Python代码与上面相同）

    ITERATIONS = 100

    python_time_per_iter = timeit(
        lambda: find_nth_prime(n), number=ITERATIONS) / ITERATIONS

    rust_time_per_iter = timeit(
        lambda: find_nth_prime_rust(n), number=ITERATIONS) / ITERATIONS

    # ...（Python代码与上面相同）

if __name__ == "__main__":
    main()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

让我们执行这段代码来比较性能。

Rust函数在寻找第N个质数时，比我们的Python实现快了5264.74%，或者说几乎是Python实现的53倍！

好了，今天又水了一篇，上面是在 Python 代码中使用 Rust 的简单示例，但有一些重要的库要么提供 Python 到 Rust 代码的绑定，要么完全用 Rust 编写。

未来AI时代，一定是注重安全性能，以及速度，一起加油了！