不少程序員都抱怨 Python/ target=_blank class=infotextkey>Python 代碼跑的慢，尤其是當處理的數據集比較大的時候。對此，本文作者指出：只需不到 100 行 Rust 代碼就能解決這個問題。
原文鏈接：https://ohadravid.Github.io/posts/2023-03-rusty-python/
作者| Ohad Ravid
譯者 | 彎月 責編 | 鄭麗媛
出品 | CSDN（ID：CSDNnews）
最近，我們的一個核心 Python 庫遇到了性能問題。這是一個非常龐大且復雜的庫，是我們 3D 處理管道的支柱，使用了 NumPy 以及其他 Python 數據科學庫來執行各種數學和幾何運算。
具體來說，我們的系統必須在 CPU 資源有限的情況下在本地運行，雖然起初的性能還不錯，但隨著并發用戶數量的增長，我們開始遇到問題，系統也出現了超負載。
我們得出的結論是：系統至少需要再快 50 倍才能處理這些增加的工作負載——而我們認為，Rust 可以幫助我們實現這一目標。
因為我們遇到的性能問題很常見，所以 下面，我來簡單介紹一下解決過程：
（a）基本的潛在問題；
（b）我們可以通過哪些優化來解決這個問題。
我們的運行示例
首先，我們通過一個小型庫來展示最初的性能問題。
假設有一個多邊形列表和一個點列表，且都是二維的，出于業務需求，我們需要將每個點“匹配”到一個多邊形。
我們的庫需要完成下列任務：
? 從點和多邊形的初始列表（全部為 2D）著手。
? 對于每個點，根據與中心的距離，找到離點最近的多邊形的子集。
? 從這些多邊形中，選擇一個“最佳”多邊形。
代碼大致如下：
fromtyping importList, Tuple importnumpy asnp fromdataclasses importdataclass fromfunctools importcached_property Point = np.array @dataclass classPolygon: x: np.array y: np.array @cached_property defcenter(self)-> Point: ... defarea(self)-> float: ... deffind_close_polygons(polygon_subset: List[Polygon], point: Point, max_dist: float)-> List[Polygon]: ... defselect_best_polygon(polygon_sets: List[Tuple[Point, List[Polygon]]])-> List[Tuple[Point, Polygon]]: ... defmain(polygons: List[Polygon], points: np.ndarray)-> List[Tuple[Point, Polygon]]: ...性能方面最主要的難點在于，Python 對象和 numpy 數組的混合。
下面，我們簡單地分析一下這個問題。
需要注意的是，對于上面這段代碼，我們當然可以把一切都轉化成 numpy 的向量計算，但真正的庫不可能這么做，因為這會導致代碼的可讀性和可修改性大大降低，收益也非常有限。此外，使用任何基于JIT 的技巧（PyPy / numba）產生的收益都非常小。
為什么不直接使用 Rust 重寫所有代碼？
雖然重寫所有代碼很誘人，但有一些問題：
? 該庫的大量計算使用了 numpy，Rust 也不一定能提高性能。
? 該庫龐大而復雜，關系到核心業務邏輯，而且高度算法化，因此重寫所有代碼需要付出幾個月的努力，而我們可憐的本地服務器眼看就要掛了。
? 一群好心的研究人員積極努力改進這個庫，實現了更好的算法，并進行了大量實驗。他們不太愿意學習一門新的編程語言，而且還要等待編譯，還要研究復雜的借用檢查器——他們不希望離開舒適區太遠。
小心探索
下面，我來介紹一下我們的分析器。
Python 有一個內置的 Profiler (cProfile)，但對于我們來說，選擇這個工具不太合適：
?它會為所有 Python 代碼引入大量開銷，卻不會給原生代碼帶來額外開銷，因此測試結果可能有偏差。
?我們將無法查看原生代碼的調用幀，這意味著我們也無法查看 Rust 代碼。
所以，我們計劃使用 py-spy，它是一個采樣分析器，可以查看原生幀。他們還將預構建的輪子發布到了 pypi，因此我們只需運行 pip install py-spy 即可。
此外，我們還需要一些測量指標。
# measure.py import time import poly_match import os# Reduce noise, actually improve perf in our case.os.environ[ "OPENBLAS_NUM_THREADS"] = "1"polygons, points = poly_match.generate_example# We are going to increase this as the code gets faster and faster.NUM_ITER = 10t0 = time.perf_counterfor _ in range(NUM_ITER):poly_match.main(polygons, points)t1 = time.perf_countertook = (t1 - t0) / NUM_ITERprint(f "Took and avg of {took * 1000:.2f}ms per iteration")
這些測量指標雖然不是很科學，但可以幫助我們優化性能。
“我們很難找到合適的測量基準。但請不要過分強調擁有完美的基準測試設置，尤其是當你優化某個程序時。”
—— Nicholas.NEThercote，《The Rust Performance Book》
運行該腳本，我們就可以獲得測量基準：
$ python measure.pyTook an avg of293.41ms per iteration對于原來的庫，我們使用了 50 個不同的樣本來確保涵蓋所有情況。
這個測量結果與實際的系統性能相符，這意味著，我們的工作就是突破這個數字。
我們還可以使用 PyPy 進行測量：
$ conda create-n pypyenv -c conda-forge pypy numpy && conda activatepypyenv $ pypy measure_with_warmup.pyTook an avgof1495.81ms per iteration先測量
首先，我們來找出什么地方如此之慢。
$py-spy record --native -o profile.svg -- python measure.py py-spy>Sampling process 100 timesa second. Press Control-C to exit. Took an avg of 365.43ms per iterationpy-spy>Stopped sampling because process exited py-spy>Wrote flamegraph data to 'profile.svg'. Samples: 391 Errors: 0我們可以看到開銷非常小。相較而言，使用 cProfile 得到的數據如下：
$ python -m cProfile measure.pyTook an avg of546.47ms per iteration 7551778functioncalls( 7409483primitive calls ) in7.806 seconds…下面是我們獲得的火焰圖：
每個方框都是一個函數，我們可以看到每個函數花費的相對時間，包括它正在調用的函數（沿著圖形/棧向下）。
要點總結：
? 絕大部分時間花在 find_close_polygons 上。
? 大部分時間都花在執行 norm，這是一個 numpy 函數。
下面，我們來仔細看看 find_close_polygons：
deffind_close_polygons(polygon_subset: List[Polygon], point: np.array, max_dist: float)-> List[Polygon]: close_polygons = []forpoly inpolygon_subset: ifnp.linalg.norm(poly.center - point) < max_dist: close_polygons.Append(poly)returnclose_polygons我們打算用 Rust 重寫這個函數。
在深入細節之前，請務必注意以下幾點：
? 此函數接受并返回復雜對象（Polygon、np.array）。
? 對象的大小非常重要（因此復制需要一定的開銷）。
? 這個函數被調用了很多次（所以我們引入的開銷可能會引發問題）。
我的第一個 Rust 模塊
PyO3 是一個用于 Python 和 Rust 之間交互的 crate ，擁有非常好的文檔。
我們將調用自己的 poly_match_rs，并添加一個名為 find_close_polygons 的函數。
mkdirpoly_match_rs && cd "$_"pipinstall maturinmaturininit --bindings pyo3maturindevelop剛開始的時候，我們的 crate 大致如下：
use pyo3::prelude::*;#[py function] fnfind_close_polygons( ) -> PyResult<( )> { Ok()}#[pymodule]fn poly_match_rs(_py: Python, m: &PyModule) -> PyResult<> {m.add_function(wrap_py function!( find_close_polygons, m)?)? ; Ok()}我們還需要記住，每次修改 Rust 庫時都需要執行 maturin develop。
改動就這么多。下面，我們來調用新函數，看看情況會怎樣。
>>> poly_match_rs.find_close_polygons( polygons, point, max_dist) ETypeError: poly_match_rs.poly_match_rs.find_close_polygonstakesnoarguments(3 given)第一版：Rust 轉換
首先，我們來定義 API。
PyO3 可以幫助我們將 Python 轉換成 Rust：
#[pyfunction]fn find_close_polygons( polygons:Vec<PyObject>, point:PyObject, max_dist:f64) -> PyResult<Vec<PyObject >> { Ok(vec![])}PyObject （顧名思義）是一個通用、“一切皆有可能”的 Python 對象。稍后，我們將嘗試與它進行交互。
這樣程序應該就可以運行了（盡管不正確）。
我直接把原來的 Python 函數復制粘帖進去，并修復了語法問題。
#[pyfunction]fn find_close_polygons( polygons: Vec<PyObject>, point: PyObject, max_dist: f64) -> PyResult<Vec<PyObject>> { letmut close_polygons = vec![];forpoly inpolygons { ifnorm( poly.center - point) < max_dist { close_polygons.push(poly)}}
Ok(close_polygons)}
可惜未能通過編譯：
% maturin develop...error[E0609]: no field `center`on type `Py<PyAny>`--> src/lib. rs:8:22|8 |ifnorm(poly.center - point) < max_dist { | ^^^^^^ unknown fielderror[E0425]: cannot find function `norm` inthis scope --> src/lib.rs:8:12|8| ifnorm(poly.center - point) < max_dist { |^^^^ notfound inthis scope error:aborting due to 2previous errors ] 58/ 59: poly_match_rs我們需要 3 個 crate 才能實現函數：
# For Rust-native array operations.ndarray= "0.15"# For a `norm` function for arrays.ndarray-linalg= "0.16"# For accessing numpy-created objects, based on `ndarray`.numpy= "0.18"首先，我們將 point: PyObject 轉換成可以使用的東西。
我們可以利用 PyO3 來轉換 numpy 數組：
use numpy::PyReadonlyArray1;#[py function] fnfind_close_polygons( // An object which says "I have the GIL", so we can access Python-managed memory.py: Python<'_>,polygons: Vec<PyObject>,// A reference to a numpy array we will be able to access.point: PyReadonlyArray1<f64>,max_dist: f64,) -> PyResult< Vec< PyObject>> { // Convert to `ndarray::ArrayView1`, a fully operational native array.letpoint = point.as_array; ...}現在 point 變成了 ArrayView1，我們可以直接使用了。例如：
// Make the `norm` function available.usendarray_linalg:: Norm; assert_eq!((point.to_owned - point).norm, 0.);接下來，我們需要獲取每個多邊形的中心，然后將其轉換成 ArrayView1。
letcenter = poly .getattr(py, "center")? // Python-style getattr, requires a GIL token (`py`)..extract::<PyReadonlyArray1<f64>>(py)? // Tell PyO3 what to convert the result to..as_array // Like `point` before..to_owned; // We need one of the sides of the `-` to be "owned".雖然信息量有點大，但總的來說，結果就是逐行轉換原來的代碼：
usepyo3::prelude::*;usendarray_linalg::Norm;usenumpy::PyReadonlyArray1;#[pyfunction]fnfind_close_polygons(py: Python<'_>,polygons: Vec<PyObject>,point: PyReadonlyArray1<f64>,max_dist: f64,)-> PyResult<Vec<PyObject>> {letmut close_polygons = vec![];letpoint = point.as_array;forpoly in polygons {letcenter = poly.getattr(py,"center")?.extract: :<PyReadonlyArray1<f64>>(py)?.as_array.to_owned;if(center - point).norm < max_dist {close_polygons.push(poly)}}Ok(close_polygons)}對比一下原來的代碼：
deffind_close_polygons(polygon_subset: List[Polygon], point: np.array, max_dist: float)-> List[Polygon]: close_polygons = []forpoly inpolygon_subset: ifnp.linalg.norm(poly.center - point) < max_dist: close_polygons.append(poly)returnclose_polygons我們希望這個版本優于原來的函數，但究竟有多少提升呢？
$( cd./poly_match_rs/ && maturin develop) $python measure.py Took an avg of 609.46ms per iteration看起來 Rust 非常慢？實則不然，使用maturin develop --release運行，就能獲得更好的結果：
$( cd./poly_match_rs/ && maturin develop --release) $python measure.py Took an avg of 23.44ms per iteration這個速度提升很不錯啊。
我們還想查看我們的原生代碼，因此發布時需要啟用調試符號。即便啟用了調試，我們也希望看到最大速度。
# added to Cargo.toml[profile.release]debug= true# Debug symbols for our profiler. lto= true# Link-time optimization. codegen-units= 1# Slower compilation but faster code.第二版：用 Rust 重寫更多代碼
接下來，在 py-spy 中通過 --native 標志，查看 Python 代碼與新版的原生代碼。
再次運行 py-spy：
$py-spy record --native -o profile.svg -- python measure.py py-spy>Sampling process 100 timesa second. Press Control-C to exit.這次得到的火焰圖如下所示（添加紅色之外的顏色，以方便參考）：
看看分析器的輸出，我們發現了一些有趣的事情：
1.find_close_polygons::...::trampoline（Python 直接調用的符號）和__pyfunction_find_close_polygons（我們的實現）的相對大小。
? 可以看到二者分別占據了樣本的 95% 和 88%，因此額外開銷非常小。
2.實際邏輯(if (center - point).norm < max_dist { ... }) 是 lib_v1.rs:22（右側非常小的框），大約占總運行時間的 9%。
? 所以應該可以實現 10 倍的提升。
3.大部分時間花在 lib_v1.rs:16 上，它是 poly.getattr(...).extract(...)，可以看到實際上只是 getattr 以及使用 as_array 獲取底層數組。
也就是說，我們需要專心解決第 3 點，而解決方法是用 Rust 重寫 Polygon。
我們來看看目標類：
@dataclassclassPolygon: x:np.array y:np.array _area:float = None @cached_propertydefcenter( self) -> np. array:centroid = np.array([ self.x, self.y]).mean(axis= 1) returncentroid defarea( self) -> float:ifself._area is None:self._area = 0. 5* np.abs( np.dot( self.x, np.roll( self.y, 1)) - np.dot( self.y, np.roll( self.x, 1)) )returnself._area我們希望盡可能保留現有的 API，但我們不需要 area 的速度大幅提升。
實際的類可能有其他復雜的東西，比如 merge 方法——使用了 scipy.spatial 中的 ConvexHull。
為了降低成本，我們只將 Polygon 的“核心”功能移至 Rust，然后從 Python 中繼承該類來實現 API 的其余部分。
我們的 struct 如下所示：
// `Array1` is a 1d array, and the `numpy` crate will play nicely with it.use ndarray::Array1;// `subclass` tells PyO3 to allow subclassing this in Python.#[pyclass(subclass)]structPolygon{ x: Array1<f64>,y: Array1<f64>,center: Array1<f64>,}下面，我們需要實現這個 struct。我們先公開 poly.{x, y, center}，作為：
? 屬性
? numpy 數組
我們還需要一個 constructor，以便 Python 創建新的 Polygon：
usenumpy::{ PyArray1, PyReadonlyArray1, ToPyArray}; #[pymethods]impl Polygon {#[new]fn new(x: PyReadonlyArray1<f64>, y: PyReadonlyArray1<f64>) -> Polygon { let x = x.as_array;let y = y.as_array;let center = Array1::from_vec(vec![x.mean.unwrap, y.mean.unwrap]);Polygon {x: x.to_owned,y: y.to_owned,center,}}// the `Py<..>` in the return type is a way of saying "an Object owned by Python".#[getter] fn x(& self, py: Python< '_>) -> PyResult<Py<PyArray1<f64>>> {Ok(self.x.to_pyarray(py).to_owned) // Create a Python-owned, numpy version of `x`.}// Same for `y` and `center`.}
我們需要將這個新的 struct 作為類添加到模塊中：
#[pymodule]fn poly_match_rs(_py: Python, m: &PyModule) -> PyResult<> {m.add_class::<Polygon>?; // new.m.add_function(wrap_py function!( find_close_polygons, m)?)? ; Ok()}然后更新 Python 代碼：
classPolygon(poly_match_rs.Polygon): _area: float = Nonedefarea(self)-> float: ...下面，編譯代碼——雖然可以運行，但速度非常慢！
為了提高性能，我們需要從 Python 的 Polygon 列表中提取基于 Rust 的 Polygon。
PyO3 可以非常靈活地處理這類操作，所以我們可以通過幾種方法來完成。我們有一個限制是我們還需要返回 Python 的 Polygon，而且我們不想克隆任何實際數據。
我們可以針對每個 PyObject 調用 .extract::<Polygon>(py)?，但也可以要求 PyO3 直接給我們 Py<Polygon>。
這是對 Python 擁有的對象的引用，我們希望它包含原生 pyclass 結構的實例（或子類，在我們的例子中）。
#[py function] fnfind_close_polygons( py: Python<'_>,polygons: Vec<Py<Polygon>>, // References to Python-owned objects.point: PyReadonlyArray1<f64>,max_dist: f64,) -> PyResult< Vec< Py< Polygon>>> { // Return the same `Py` references, unmodified.letmut close_polygons = vec![]; letpoint = point.as_array; forpoly inpolygons { letcenter = poly.borrow(py).center // Need to use the GIL (`py`) to borrow the underlying `Polygon`..to_owned;if(center - point).norm < max_dist { close_polygons.push(poly)}}Ok(close_polygons)}下面，我們來看看使用這些代碼的效果如何：
$ python measure.pyTook an avg of6.29ms per iteration我們快要成功了，只需再提升一倍的速度即可。
第三版：避免內存分配
我們再來看一看分析器的結果。
1.首先，我們看到 select_best_polygon，現在它調用的是一些 Rust 代碼（在獲取 x 和 y 向量時）。
? 我們可以解決這個問題，但這是一個非常小的提升（大約為 10%）。
2.我們看到 extract_argument 花費了大約 20% 的時間（在 lib_v2.rs:48 下），這個開銷相對比較大。
? 但大部分時間都花在了 PyIterator::next 和 PyTypeInfo::is_type_of 中，這可不容易修復。
3.我們看到大量時間花在了內存分配上。
? lib_v2.rs:58 是我們的 if 語句，我們還看到了drop_in_place和to_owned。
? 實際的代碼大約占總時間的 35%，遠超我們的預期。所有數據都已存在，所以這一段本應非常快。
下面，我們來解決最后一點。
有問題的代碼如下：
letcenter = poly.borrow(py).center .to_owned;if(center - point).norm < max_dist { ... }我們希望避免 to_owned。但是，我們需要一個已擁有的 norm 對象，所以我們必須手動實現。
具體的寫法如下：
use ndarray_linalg::Scalar;let center = &poly.as_ref(py).borrow.center;if((center[ 0] - point[ 0]).square + (center[ 1] - point[ 1]).square). sqrt< max_dist { close_polygons.push(poly)}然而，借用檢查器報錯了：
error[E0505]: cannot move out of `poly`because it is borrowed --> src/lib. rs:58:33|55 |let center = &poly.as_ref(py).borrow.center; | ------------------------|||borrow of `poly`occurs here | a temporary with access to the borrow is created here ......58 |close_polygons.push(poly); | ^^^^ move out of `poly` occurs here59 |} 60| }|- ... andthe borrow might be used here, whenthat temporary is dropped andruns the `Drop`code fortype `PyRef`借用檢查器是正確的，我們使用內存的方式不正確。
更簡單的修復方法是直接克隆，然后 close_polygons.push(poly.clone) 就可以通過編譯了。
這實際上是一個開銷很低的克隆，因為我們只增加了 Python 對象的引用計數。
然而，在這個例子中，我們也可以通過一個 Rust 的常用技巧：
letnorm = { letcenter = &poly.as_ref(py).borrow.center; ((center[ 0] - point[ 0]).square + (center[ 1] - point[ 1]).square).sqrt };ifnorm < max_dist { close_polygons.push(poly)}由于 poly 只在內部范圍內被借用，如果我們接近 close_polygons.pus，編譯器就可以知道我們不再持有引用，因此就可以通過編譯。
最后的結果：
$ python measure.pyTook an avg of2.90ms per iteration相較于原來的代碼，整體性能得到了 100 倍的提升。
總結
我們原來的 Python 代碼如下：
@dataclassclassPolygon: x: np.arrayy: np.array_area: float = None@cached_propertydefcenter(self)-> np.array: centroid = np.array([self.x, self.y]).mean(axis= 1) returncentroid defarea(self)-> float: ...deffind_close_polygons(polygon_subset: List[Polygon], point: np.array, max_dist: float)-> List[Polygon]: close_polygons = []forpoly inpolygon_subset: ifnp.linalg.norm(poly.center - point) < max_dist: close_polygons.append(poly)returnclose_polygons # Rest of file (main, select_best_polygon).我們使用 py-spy 對其進行了分析，即便用最簡單的、逐行轉換的 find_close_polygons，也可以獲得 10 倍的性能提升。
我們反復進行分析-修改代碼-測量結果，并最終獲得了 100 倍的性能提升，同時 API 仍然保持與原來的庫相同。
最終得到的 Python 代碼如下：
importpoly_match_rs frompoly_match_rs importfind_close_polygons classPolygon(poly_match_rs.Polygon): _area: float = Nonedefarea(self)-> float: ...# Rest of file unchanged (main, select_best_polygon).調用的 Rust 代碼如下：
usepyo3::prelude::*;usendarray::Array1;usendarray_linalg::Scalar;usenumpy::{PyArray1, PyReadonlyArray1, ToPyArray};#[pyclass(subclass)]structPolygon {x: Array1<f64>,y: Array1<f64>,center: Array1<f64>,}#[pymethods]implPolygon {#[new]fnnew(x: PyReadonlyArray1<f64>, y: PyReadonlyArray1<f64>) -> Polygon {letx = x.as_array;lety = y.as_array;letcenter = Array1::from_vec(vec![x.mean.unwrap, y.mean.unwrap]);Polygon{x: x.to_owned,y: y.to_owned,center,}}#[getter]fnx(&self, py: Python<'_>) -> PyResult<Py<PyArray1<f64>>> {Ok(self.x.to_pyarray(py).to_owned)}//Same for `y` and `center`.}#[pyfunction]fnfind_close_polygons(py: Python<'_>,polygons: Vec<Py<Polygon>>,point: PyReadonlyArray1<f64>,max_dist: f64,)-> PyResult<Vec<Py<Polygon>>> {letmut close_polygons = vec![];letpoint = point.as_array;forpoly in polygons {letnorm = {letcenter = &poly.as_ref(py).borrow.center;((center[0]- point[0]).square + (center[1] - point[1]).square).sqrt};ifnorm < max_dist {close_polygons.push(poly)}}Ok(close_polygons)}#[pymodule]fnpoly_match_rs(_py: Python, m: &PyModule) -> PyResult<> {m.add_class: :<Polygon>?;m.add_function(wrap_pyfunction!(find_close_polygons,m)?)?;Ok()}要點總結
? Rust（在 PyO3 的幫助下）能夠以非常小的代價換取 Python 代碼性能的大幅提升。
? 對于研究人員來說，Python API 非常優秀，同時使用 Rust 快速構建基本功能是一個非常強大的組合。
? 分析非常有趣，可以幫助你了解代碼中的一切。
?最希望ChatGPT開源，一半開發者參與過開源貢獻，63%的人在用愛發電|中國開源開發者現狀
? 微軟總裁稱中國將是 ChatGPT 主要對手；曝蘋果 M3 芯片下半年量產；linux 6.3 正式發布|極客頭條
? ChatGPT 搶不走程 序員飯碗的原因找到了？最新研究：它自動生成了 21 個程序，16 個有漏 洞
