PyO3传递Lambda函数作为参数的性能瓶颈及优化咨询

性能差距是否正常？ #

这种近两个数量级的性能差距是正常现象。核心原因是：原Rust版本的回调是纯Rust闭包，函数调用是编译期优化后的零开销操作；而通过PyO3调用Python Lambda时，每次回调都要经历以下不可避免的开销：

• GIL的获取与释放（PyO3自动处理，但仍有固定开销）
• Python解释器的函数调用栈帧创建、参数打包与销毁
• Python对象到Rust f64的类型转换（extract方法包含通用类型分支检查逻辑）

你的场景中回调被执行了22.5万次，这些单次微小的开销会被持续放大，最终导致总耗时飙升至20ms级别。

无需修改原Rust test_function的优化方案 #

1. 用PyFunction替代PyAny，减少类型检查开销 #

将封装函数的参数类型从&PyAny改为&PyFunction，可以跳过is_callable()的运行时检查，直接调用函数，减少不必要的开销：

#[pyfunction]
pub fn py_test_function(function: &PyFunction) -> f64 {
    let cb = move |x| function.call1((x,)).unwrap().extract::<f64>().unwrap();
    test_function(cb)
}

2. 直接downcast到PyFloat，优化类型转换 #

extract::<f64>()是通用类型转换方法，包含大量类型分支检查。可以直接将调用结果downcast为PyFloat，再通过value()方法直接获取原始f64值，大幅减少类型转换开销：

use pyo3::types::PyFloat;

#[pyfunction]
pub fn py_test_function(function: &PyFunction) -> f64 {
    let cb = move |x| {
        let py_result = function.call1((x,)).unwrap();
        py_result.downcast::<PyFloat>().unwrap().value()
    };
    test_function(cb)
}

针对你更新方案的进一步优化 #

你尝试的批量处理思路是正确的，但耗时未降低的核心原因是numpy数组到Rust Vec<f64>的内存拷贝开销。可以直接借用numpy数组的内存，避免复制：

修改Rust封装函数，直接使用PyArray接收numpy数组 #

use pyo3::types::PyArray1;

#[pyfunction]
pub fn py_test_function_alt(values: &PyArray1<f64>) -> f64 {
    // 直接获取numpy数组的切片，无需内存拷贝
    let slice = values.as_slice().unwrap();
    test_function_alt(slice)
}

对应的Python调用代码保持不变 #

import numpy as np
import pyo3test

def foobar(cb):
    vals = cb(np.arange(225000))
    return pyo3test.py_test_function_alt(vals)

这种方式下，Rust直接读取numpy数组的原始内存，避免了整个数组的拷贝操作，能显著降低耗时。

内容的提问来源于stack exchange，提问作者Quasark