Ускоряем Python в сто раз при помощи менее чем ста строк на Rust

Однажды на работе у нас возникла проблема с производительностью одной из наших основных Python-библиотек.

Эта библиотека формирует фундамент нашего конвейера 3D-обработки. Это довольно большая и сложная библиотека, использующая NumPy и другие научные пакеты Python для выполнения широкого спектра математических и геометрических операций.

Кроме того, наша система должна работать на мощностях компании с ограниченными ресурсами CPU, и хотя поначалу она справлялась хорошо, с ростом количества одновременных физических пользователей у нас начали возникать проблемы, а наша система едва выдерживала нагрузку.

Мы пришли к выводу, что для работы с увеличившейся нагрузкой нужно сделать систему, по крайней мере, в пятьдесят раз быстрее, и решили, что помочь в этом нам может Rust.

Так как проблемы производительности, с которыми мы столкнулись, встречаются довольно часто, мы можем воссоздать и решить их прямо здесь, в (не такой уж короткой) статье.

Так что заварите себе чаю (или кофе) и я расскажу вам (а) о самой проблеме и (б) о нескольких итерациях оптимизаций, которые мы смогли применить для решения этой проблемы.

Если вы хотите сразу перейти к получившемуся коду, то читайте раздел «Подведение итогов».

▍ Наш пример

Давайте создадим небольшую библиотеку, которая будет демонстрировать наши исходные проблемы производительности (однако выполняющая совершенно произвольную работу).

Представьте, что у вас есть список многоугольников и список точек, всё это в 2D. Для бизнес-логики нам нужно «сопоставить» каждую точку с одним многоугольником.

Наша вымышленная библиотека будет:

1. Начинать с исходного списка точек и многоугольников (в 2D).
2. Для каждой точки на основании расстояния до центра находить гораздо меньшее подмножество многоугольников, ближайших к ней.
3. Из этих многоугольников выбирать «лучший» (в качестве показателя «лучшего» мы будем использовать «с наименьшей площадью»).

В коде это будет выглядеть так (полный код можно найти здесь):

from typing import List, Tuple
import numpy as np
from dataclasses import dataclass
from functools import cached_property

Point = np.array

@dataclass
class Polygon:
    x: np.array
    y: np.array

    @cached_property
    def center(self) -> Point: ...
    def area(self) -> float: ...

def find_close_polygons(polygon_subset: List[Polygon], point: Point, max_dist: float) -> List[Polygon]:
    ...

def select_best_polygon(polygon_sets: List[Tuple[Point, List[Polygon]]]) -> List[Tuple[Point, Polygon]]:
    ...

def main(polygons: List[Polygon], points: np.ndarray) -> List[Tuple[Point, Polygon]]:
    ...

Основная сложность (с точки зрения производительности) заключается в смешении объектов Python и массивов numpy.

Скоро мы подробно это проанализируем.

Стоит заметить, что для этого искусственного примера библиотеки можно преобразовать части/всё целиком в векторизированный numpy, но это практически невозможно для реальной библиотеки и при этом делает код гораздо менее читаемым и модифицируемым, а преимущества окажутся незначительными (здесь находится частично векторизированная версия, которая быстрее, но далека от результатов, которые мы достигнем).

Кроме того, использование любых трюков с JIT (PyPy / numba) приводит к очень маленькой выгоде (чтобы убедиться, мы измерим это).

▍ Почему бы просто не Переписать Всё На Rust™?

Как ни привлекательно было полное переписывание, оно имело несколько проблем:

1. Библиотека уже использовала numpy для многих своих вычислений, так почему мы должны ожидать, что Rust будет лучше?
2. Она большая, сложная, очень важная для бизнеса и высокоалгоритмическая, поэтому это займёт месяцы работы, а бедный сервер компании помирает уже сейчас.
3. Группа дружественных исследователей активно работала над этой библиотекой, реализуя улучшенные алгоритмы и проводя множество экспериментов. Их не очень порадует необходимость изучения нового языка программирования, ожидания компиляции и борьбы с borrow checker. Они будут признательны, если мы не станем так усложнять их работу.

▍ Экспериментируем

Настало время добавить нашего друга — профилировщик.

Python имеет встроенный профилировщик (cProfile), но в данном случае это неподходящий инструмент для работы:

1. Он навесит большую трату ресурсов на весь код на Python, но не увеличит нагрузку на нативный код, поэтому результаты могут оказаться искажёнными.
2. Мы не сможем просматривать нативные кадры, то есть не сможем заглянуть в наш код на Rust.

Мы будем использовать py-spy (GitHub).

py-spy — это сэмплирующий профилировщик, способный заглядывать в нативные кадры.

Разработчики также любезно опубликовали готовые wheel для pypi, поэтому можно просто ввести pip install py-spy и приступать к работе.

Ещё нам нужно что-то для измерений.

# measure.py
import time
import poly_match
import os
 
# Снижаем шум, в нашем случае повышая производительность.
os.environ["OPENBLAS_NUM_THREADS"] = "1"

polygons, points = poly_match.generate_example()

# Мы будем увеличивать это значение, когда наш код будет становиться всё быстрее и быстрее.
NUM_ITER = 10

t0 = time.perf_counter()
for _ in range(NUM_ITER):
    poly_match.main(polygons, points)
t1 = time.perf_counter()

took = (t1 - t0) / NUM_ITER
print(f"Took and avg of {took * 1000:.2f}ms per iteration")

Выглядит не очень научно, но позволит нам сделать огромный шаг вперёд.

«Хороший бенчмаркинг — это сложно. Тем не менее — не слишком налегайте на создание идеальной системы бенчмаркинга, особенно когда вы начинаете оптимизировать программу». ~ Николас Нетеркот, «The Rust Performance Book»

После выполнения этого скрипта мы получим результаты, от которых можно отталкиваться:

$ python measure.py
Took an avg of 293.41ms per iteration

Для нашей настоящей библиотеки мы использовали 50 разных случаев, чтобы гарантировать покрытие всех случаев.

Это соответствовало общей производительности системы, то есть мы могли приступать к борьбе с этими числами.

Примечание: можно также выполнять измерения при помощи PyPy (также мы добавим разогрев, чтобы JIT показала свою магию).

$ conda create -n pypyenv -c conda-forge pypy numpy && conda activate pypyenv
$ pypy measure_with_warmup.py
Took an avg of 1495.81ms per iteration

▍ Измерения — это главное

Так, давайте выясним, что же тут так медленно работает.

$ py-spy record --native -o profile.svg -- python measure.py
py-spy> Sampling process 100 times a second. Press Control-C to exit.

Took an avg of 365.43ms per iteration

py-spy> Stopped sampling because process exited
py-spy> Wrote flamegraph data to 'profile.svg'. Samples: 391 Errors: 0

Мы уже видим, что лишняя трата ресурсов тут довольно мала. Просто для сравнения, при использовании cProfile мы получим следующее:

$ python -m cProfile measure.py
Took an avg of 546.47ms per iteration
         7551778 function calls (7409483 primitive calls) in 7.806 seconds
         ...

Мы получаем вот такой красивый красный график, называемый flamegraph:

Каждый прямоугольник — это функция. В оригинале статьи график интерактивен и мы можем видеть относительное время, проводимое в каждой функции, в том числе в функциях, которые она вызывает (спускаясь вниз по графику/стеку). Попробуйте нажать на прямоугольник norm, чтобы приблизить его.

Здесь можно сделать следующие основные выводы:

1. Подавляющее большинство времени тратится в find_close_polygons.
2. Основная часть этого времени тратится на выполнение norm, которая является функцией numpy.

Итак, давайте взглянем на find_close_polygons:

def find_close_polygons(
    polygon_subset: List[Polygon], point: np.array, max_dist: float
) -> List[Polygon]:
    close_polygons = []
    for poly in polygon_subset:
        if np.linalg.norm(poly.center - point) < max_dist:
            close_polygons.append(poly)

    return close_polygons

Мы перепишем эту функцию на Rust.

Прежде чем вдаваться в подробности, важно заметить здесь несколько моментов:

1. Эта функция принимает и возвращает сложные объекты (Polygon, np.array).
2. Размер объектов нетривиален (поэтому копирование может быть затратным).
3. Эта функция вызывается часто (поэтому внесённая нами лишняя трата ресурсов, вероятно, будет иметь значение).

▍ Мой первый модуль Rust

pyo3 — это крейт для взаимодействия между Python и Rust. Он имеет чрезвычайно хорошую документацию, а базовая настройка объяснена здесь.

Мы назовём свой крейт poly_match_rs и добавим функцию find_close_polygons.

mkdir poly_match_rs && cd "$_"
pip install maturin
maturin init --bindings pyo3
maturin develop

Изначально наш крейт будет выглядеть так:

use pyo3::prelude::*;

#[pyfunction]
fn find_close_polygons() -> PyResult<()> {
    Ok(())
}

#[pymodule]
fn poly_match_rs(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(find_close_polygons, m)?)?;
    Ok(())
}

Нам также нужно не забывать выполнять maturin develop каждый раз, когда мы меняем библиотеку Rust.

Вот и всё! Давайте вызовем нашу новую функцию и посмотрим, что произойдёт.

>>> poly_match_rs.find_close_polygons(polygons, point, max_dist)
E TypeError: poly_match_rs.poly_match_rs.find_close_polygons() takes no arguments (3 given)

▍ v1 — наивная трансляция 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>> {
    let mut close_polygons = vec![];
    
    for poly in polygons {
        if norm(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 |         if norm(poly.center - point) < max_dist {
  |                      ^^^^^^ unknown field


error[E0425]: cannot find function `norm` in this scope
 --> src/lib.rs:8:12
  |
8 |         if norm(poly.center - point) < max_dist {
  |            ^^^^ not found in this scope


error: aborting due to 2 previous errors ] 58/59: poly_match_rs

Для реализации этой функции нам потребуются три крейта:

# Для нативных операций с массивами Rust.
ndarray = "0.15"

# Для функции "norm" для массивов.
ndarray-linalg = "0.16"  

# Для доступа к созданным numpy объектам, основанным на "ndarray".
numpy = "0.18"

Для начала давайте превратим непрозрачную и обобщённую point: PyObject во что-то, с чем мы сможем работать.

Точно так же, как мы попросили у PyO3 Vec объекта PyObject, мы можем попросить numpy-array, и он автоматически преобразует аргумент.

use numpy::PyReadonlyArray1;

#[pyfunction]
fn find_close_polygons(
    // Объект, говорящий "У меня есть GIL", чтобы мы могли получить доступ к управляемой Python памяти.
    py: Python<'_>,
    polygons: Vec<PyObject>,
    // Ссылка на массив numpy, доступ к которой мы сможем иметь.
    point: PyReadonlyArray1<f64>,
    max_dist: f64,
) -> PyResult<Vec<PyObject>> {
    // Преобразуем в полнофункциональный нативный массив "ndarray::ArrayView1".
    let point = point.as_array();
    ...
}

Так как точка теперь находится в ArrayView1, мы можем её использовать. Например:

// делаем доступной функцию "norm".
use ndarray_linalg::Norm;

assert_eq!((point.to_owned() - point).norm(), 0.);

Теперь нам просто нужно получить центр каждого многоугольника и преобразовать его в ArrayView1.

На PyO3 это выглядит так:

let center = poly
  .getattr(py, "center")?                 // getattr в стиле Python, требует токен GIL (`py`).
  .extract::<PyReadonlyArray1<f64>>(py)?  // Сообщаем PyO3, во что преобразовать результат.
  .as_array()                             // Как "point" ранее.
  .to_owned();                            // Нам нужно, чтобы одной из сторон `-` "владели".

Слегка многословно, но в целом результат, очевидно, является построчной трансляцией первоначального кода:

use pyo3::prelude::*;

use ndarray_linalg::Norm;
use numpy::PyReadonlyArray1;

#[pyfunction]
fn find_close_polygons(
    py: Python<'_>,
    polygons: Vec<PyObject>,
    point: PyReadonlyArray1<f64>,
    max_dist: f64,
) -> PyResult<Vec<PyObject>> {
    let mut close_polygons = vec![];
    let point = point.as_array();
    for poly in polygons {
        let center = 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)
}

Сравним с оригиналом:

def find_close_polygons(
    polygon_subset: List[Polygon], point: np.array, max_dist: float
) -> List[Polygon]:
    close_polygons = []
    for poly in polygon_subset:
        if np.linalg.norm(poly.center - point) < max_dist:
            close_polygons.append(poly)

    return close_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       # Отладочные символы для нашего профилировщика.
lto = true         # Оптимизация этапа компоновки.
codegen-units = 1  # Более медленная компиляция, зато более быстрый код. 

▍ v2 — перепишем ещё больше на Rust

Теперь использовав флаг --native в py-spy, мы сможем видеть и Python, и наш новый нативный код.

Снова запускаем py-spy:

$ py-spy record --native -o profile.svg -- python measure.py
py-spy> Sampling process 100 times a second. Press Control-C to exit.

Получаем такой flamegraph (некрасные цвета добавлены для того, чтобы мы могли отталкиваться от них):

Интерактивную версию можно посмотреть здесь

Посмотрев на вывод профилировщика, мы можем заметить несколько интересных моментов:

1. Относительный размер find_close_polygons::...::trampoline (символа, который Python непосредственно вызывает) и __pyfunction_find_close_polygons (нашей реализации).
   • Если навести мышь, видно, что сэмплы соотносятся как 95% и 88%, то есть лишняя трата ресурсов довольно мала.
2. Сама логика (if (center - point).norm() < max_dist { ... }), которая находится в lib_v1.rs:22 (очень маленький прямоугольник справа), занимает примерно 9% от всего времени исполнения.
   • То есть улучшение в десять раз всё равно должно быть возможно!
3. Основная часть времени тратится на lib_v1.rs:16, то есть на poly.getattr(...).extract(...), и если мы приблизим изображение, то увидим, что на самом деле это просто getattr и получение внутреннего массива при помощи as_array.

Можно сделать вывод, что нам нужно сосредоточиться на решении третьего пункта, и сделать это можно Переписав Polygon на Rust.

Давайте взглянем на нашу цель:

@dataclass
class Polygon:
    x: np.array
    y: np.array
    _area: float = None

    @cached_property
    def center(self) -> np.array:
        centroid = np.array([self.x, self.y]).mean(axis=1)
        return centroid

    def area(self) -> float:
        if self._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))
            )
        return self._area

Нам нужно максимально сохранить имеющийся API, но нам не нужно (пока), чтобы area была такой уж быстрой.

Сам класс может содержать дополнительные сложные вещи, например, метод merge, использующий ConvexHull из scipy.spatial.

Чтобы снизить затраты (и ограничить рамки этой и так длинной статьи), мы только переместим «основную» функциональность Polygon на Rust, и сделаем её подклассом из Python для реализации остального API.

Наша struct будет выглядеть так:

// "Array1" - это одномерный массив, и крейт "numpy" будет хорошо с ним взаимодействовать.
use ndarray::Array1;

// "subclass" приказывает PyO3 разрешить создание подклассов этого на Python.
#[pyclass(subclass)]
struct Polygon {
    x: Array1<f64>,
    y: Array1<f64>,
    center: Array1<f64>,
}

Мы хотим раскрыть poly.{x, y, center} следующим образом:

1. Свойства.
2. Массивы numpy.

Также нам понадобится конструктор, чтобы Python мог создавать новые Polygon.

use numpy::{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,
        }
    }
    
    // "Py<..>" в возвращаемом типе - это способ показать, что объект принадлежит Python".
    #[getter]               
    fn x(&self, py: Python<'_>) -> PyResult<Py<PyArray1<f64>>> {
        Ok(self.x.to_pyarray(py).to_owned()) // Создаём numpy-версию "x", которой будет владеть Python.
    }

    // То же для "y" и "center".
}

Мы должны добавить нашу новую struct как класс в модуль:

#[pymodule]
fn poly_match_rs(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_class::<Polygon>()?; // new.
    m.add_function(wrap_pyfunction!(find_close_polygons, m)?)?;
    Ok(())
}

А теперь мы можем дополнить код на Python, чтобы использовать его:

class Polygon(poly_match_rs.Polygon):
    _area: float = None

    def area(self) -> float:
        ...

Можно выполнить компиляцию, и это на самом деле заработает, только гораздо медленнее! (Помните, что x, y и center теперь нужно будет создавать новый массив numpy при каждой операции доступа).

Чтобы повысить производительность, нам нужно извлечь исходный Polygon на Rust из списка Polygon Python.

PyO3 очень гибок с таким типом операций, поэтому мы можем сделать это несколькими способами. Ограничение заключается в том, что нам также нужно возвращать Polygon языка Python и мы не хотим выполнять клонирование данных.

Можно вручную вызывать .extract::<Polygon>(py)? для каждого PyObject, но мы попросим PyO3 давать нам непосредственно Py<Polygon>.

Это ссылка на принадлежащий Python объект, который, как мы ожидаем, содержит экземпляр (или в нашем случае подкласс) нативной структуры pyclass.

#[pyfunction]
fn find_close_polygons(
    py: Python<'_>,
    polygons: Vec<Py<Polygon>>,             // Ссылки на объекты, которыми владеет Python.
    point: PyReadonlyArray1<f64>,
    max_dist: f64,
) -> PyResult<Vec<Py<Polygon>>> {           // Возвращает те же ссылки на "Py" без изменений.
    let mut close_polygons = vec![];
    let point = point.as_array();
    for poly in polygons {
        let center = poly.borrow(py).center // Должен использовать GIL ("py"), чтобы позаимствовать внутренний "Polygon".
            .to_owned();

        if (center - point).norm() < max_dist {
            close_polygons.push(poly)
        }
    }

    Ok(close_polygons)
}

Посмотрим, чего мы добились при помощи этого кода:

$ python measure.py
Took an avg of 6.29ms per iteration

Мы почти у цели! Осталось удвоить производительность!

▍ v3 — избегаем распределений

Давайте ещё раз запустим профилировщик.

Интерактивная версия находится здесь

1. Мы начинаем видеть select_best_polygon, которая теперь вызывает какой-то код на Rust (когда получает векторы x и y)
   • Можно это исправить, но потенциал этого улучшения очень мал (возможно, 10%)
2. Мы видим, что тратим примерно 20% времени в extract_argument (в lib_v2.rs:48), то есть лишняя трата ресурсов по-прежнему довольно велика!
   • Но бОльшая часть времени тратится на PyIterator::next и PyTypeInfo::is_type_of, что исправить не так просто.
3. Мы видим, что довольно много времени тратится на распределение!
   • lib_v2.rs:58 — это наш if, и мы видим drop_in_place и to_owned.
   • Всего это около 35% от суммарного времени, что гораздо больше, чем мы ожидали: при наличии всех данных это должна быть быстрая часть кода.

Давайте разберёмся с последним пунктом.

Вот фрагмент кода, вызывающий проблемы:

let center = poly.borrow(py).center
    .to_owned();

if (center - point).norm() < max_dist { ... } 

Мы хотим избежать вот этого to_owned. Но нам нужен owned-объект для norm, поэтому придётся реализовать его вручную.

(Причина того, что мы можем улучшить ndarray, заключается в том, что мы знаем, что наш массив на самом деле представляет собой два значения f32).

Это будет выглядеть так:

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)
}

Но, увы, borrow checker нами недоволен:

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 here
59 |         }
60 |     }
   |     - ... and the borrow might be used here, when that temporary is dropped and runs the `Drop` code for type `PyRef`

Как обычно, borrow checker прав: мы совершаем преступления с памятью.

Проще всего исправить это при помощи Just Clone, после чего close_polygons.push(poly.clone()) скомпилируется.

На самом деле, это очень малозатратный клон, потому что мы всего лишь выполняем incr счётчика ссылок объекта Python.

Однако в этом случае мы можем сократить borrow, сделав классический трюк Rust:

let norm = {
    let center = &poly.as_ref(py).borrow().center;

    ((center[0] - point[0]).square() + (center[1] - point[1]).square()).sqrt()
};

if norm < max_dist {
    close_polygons.push(poly)
}

Так как poly заимствуется только во внутренней области видимости, после того как мы достигнем close_polygons.push, компилятор сможет узнать, что мы больше не храним эту ссылку и спокойно скомпилирует новую версию.

Наконец, мы добились следующего:

$ python measure.py
Took an avg of 2.90ms per iteration

Стократное улучшение по сравнению с первоначальным кодом.

▍ Подведение итогов

Мы начинали с такого кода на Python:

@dataclass
class Polygon:
    x: np.array
    y: np.array
    _area: float = None

    @cached_property
    def center(self) -> np.array:
        centroid = np.array([self.x, self.y]).mean(axis=1)
        return centroid

    def area(self) -> float:
        ...

def find_close_polygons(
    polygon_subset: List[Polygon], point: np.array, max_dist: float
) -> List[Polygon]:
    close_polygons = []
    for poly in polygon_subset:
        if np.linalg.norm(poly.center - point) < max_dist:
            close_polygons.append(poly)

    return close_polygons

# Остальная часть файла (main, select_best_polygon).

Мы профилировали его при помощи py-spy, и даже самая наивная построчная трансляция find_close_polygons обеспечила улучшение в десять с лишним раз.

Мы выполнили несколько дополнительных итераций «профилирование-написание кода-измерения», пока не получили стократное улучшение времени исполнения, сохранив при этом тот же API, что и у исходной библиотеки.

Версия	Ср. время на итерацию (мс)	Множитель
Базовая реализация (на Python)	293,41	1x
Наивная построчная трансляция на Rust find_close_polygons	23,44	12,50x
Реализация Polygon на Rust	6,29	46,53x
Оптимизированная реализация распределений на Rust	2,90	101,16x

Окончательный код на Python выглядит так:

import poly_match_rs
from poly_match_rs import find_close_polygons

class Polygon(poly_match_rs.Polygon):
    _area: float = None

    def area(self) -> float:
        ...

# Оставшаяся часть файла не изменилась (main, select_best_polygon).

Он вызывает следующий код на Rust:

use pyo3::prelude::*;

use ndarray::Array1;
use ndarray_linalg::Scalar;
use numpy::{PyArray1, PyReadonlyArray1, ToPyArray};

#[pyclass(subclass)]
struct Polygon {
    x: Array1<f64>,
    y: Array1<f64>,
    center: Array1<f64>,
}

#[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,
        }
    }

    #[getter]
    fn x(&self, py: Python<'_>) -> PyResult<Py<PyArray1<f64>>> {
        Ok(self.x.to_pyarray(py).to_owned())
    }

    // То же для "y" и "center".
}

#[pyfunction]
fn find_close_polygons(
    py: Python<'_>,
    polygons: Vec<Py<Polygon>>,
    point: PyReadonlyArray1<f64>,
    max_dist: f64,
) -> PyResult<Vec<Py<Polygon>>> {
    let mut close_polygons = vec![];
    let point = point.as_array();
    for poly in polygons {
        let norm = {
            let center = &poly.as_ref(py).borrow().center;

            ((center[0] - point[0]).square() + (center[1] - point[1]).square()).sqrt()
        };

        if norm < max_dist {
            close_polygons.push(poly)
        }
    }

    Ok(close_polygons)
}

#[pymodule]
fn poly_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 — это чрезвычайно мощное сочетание.
• Профилирование крайне интересно, оно мотивирует по-настоящему разобраться со всем, что происходит в вашем коде.

И последнее: компьютеры безумно быстры. Когда вам в следующий раз нужно будет что-то сделать, попробуйте запустить профилировщик, возможно, вы научитесь чему-то новому.

Пол-лимона подарков от RUVDS. Отвечай на вопросы и получай призы 🍋