Использование плагинов Polars для повышения скорости в 14 раз с помощью Rust

Увеличение скорости в 14 раз с помощью плагинов Polars и языка программирования Rust

Достижение высокой скорости вне собственной библиотеки Polars

Введение

Polars завоевывает мир благодаря своей скорости, эффективности использования памяти и красивому интерфейсу программирования. Если вы хотите узнать, насколько мощная она, посмотрите на бенчмарки DuckDB. Причем это еще не с использованием самой последней версии Polars.

Несмотря на все удивительные возможности, которые Polars может предложить, она не всегда является лучшим решением для выполнения ВСЕХ вычислений, которые вам могут понадобиться. Существуют несколько исключений, в которых Polars не превосходит другие инструменты. Однако с недавним выпуском системы плагинов Polars для Rust, это может уже не быть так.

Плагины Polars

Что такое плагин Polars? Это просто способ создания собственных выражений Polars с использованием нативного Rust и предоставления доступа к этим выражениям через специальное пространство имен. Это позволяет использовать скорость Rust и применять ее к вашему DataFrame в Polars для выполнения вычислений с использованием скорости и встроенных инструментов, предоставляемых Polars.

Давайте рассмотрим несколько конкретных примеров.

• Конфиденциальность данных и ее влияние на управление
• 7 примеров для освоения операций с категориальными данными с помощью Python Pandas
• AI Кодинг Является ли Google Bard хорошим разработчиком Python?

Последовательные вычисления

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

У нас есть следующий алгоритм для вычисления кумулятивного значения массива чисел для данного ряда, определенного как набор чисел с одинаковым знаком. Например:

┌───────┬───────────┐│ value ┆ run_value ││ ---   ┆ ---       ││ i64   ┆ i64       │╞═══════╪═══════════╡│ 1     ┆ 1         │   # Первый ряд начинается здесь│ 2     ┆ 3         ││ 3     ┆ 6         ││ -1    ┆ -1        │   # Ряд сбрасывается здесь│ -2    ┆ -3        ││ 1     ┆ 1         │   # Ряд сбрасывается здесь└───────┴───────────┘

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

Начнем с базовой версии, написанной на pandas.

def calculate_runs_pd(s: pd.Series) -> pd.Series:    out = []    is_positive = True    current_value = 0.0    for value in s:        if value > 0:            if is_positive:                current_value += value            else:                current_value = value                is_positive = True        else:            if is_positive:                current_value = value                is_positive = False            else:                current_value += value        out.append(current_value)    return pd.Series(out)

Мы перебираем серию, вычисляем текущее значение для каждой позиции и возвращаем новую серию Pandas.

Время выполнения

Прежде чем продолжить, давайте проведем несколько тестов. Мы измерим скорость выполнения и потребление памяти с помощью pytest-benchmark и pytest-memray. Мы настроим проблему таким образом, что у нас будет столбец сущностей, временной столбец и столбец функций. Задача состоит в том, чтобы вычислить значения рядов для каждой сущности в данных по времени. Мы установим количество сущностей и временных меток каждое равным 1 000, что даст нам DataFrame с 1 000 000 строками.

При запуске нашей реализации Pandas против замеров с использованием функциональности groupby apply в Pandas, мы получаем следующие результаты:

Наивная реализация Polars

Итак, у нас есть наши замеры. Давайте посмотрим, как реализовать эту же функциональность в Polars. Мы начнем с очень похожей версии, которая будет применена путем применения функции к объекту GroupBy в Polars.

def calculate_runs_pl_apply(s: pl.Series) -> pl.DataFrame:    out = []    is_positive = True    current_value = 0.0    for value in s:        if value is None:            pass        elif value > 0:            if is_positive:                current_value += value            else:                current_value = value                is_positive = True        else:            if is_positive:                current_value = value                is_positive = False            else:                current_value += value        out.append(current_value)    return pl.DataFrame(pl.Series("run", out))

Теперь давайте посмотрим, как это соотносится с нашим исходным замером в Pandas.

Что-то не очень хорошо получилось. Это не должно быть сюрпризом. Авторы Polars очень ясно указали, что очень распространенный подход группировки apply в Pandas неэффективен для выполнения вычислений в Polars. Здесь это показывается. Как скорость, так и потребление памяти хуже, чем у нашей исходной реализации в Pandas.

Реализация выражений Polars

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

def calculate_runs_pl_native(df: pl.LazyFrame, col: str, by: str) -> pl.LazyFrame:    return (        df.with_columns((pl.col(col) > 0).alias("__is_positive"))        .with_columns(            (pl.col("__is_positive") != pl.col("__is_positive").shift(1))            .over(by)            .fill_null(False)            .alias("__change_sides")        )        .with_columns(pl.col("__change_sides").cumsum().over(by).alias("__run_groups"))        .with_columns(pl.col(col).cumsum().over(by, "__run_groups").alias("runs"))        .select(~cs.starts_with("__"))    )

Краткое объяснение того, что мы делаем здесь:

• Найти все строки, где признак положительный
• Найти все строки, где столбец __is_positive отличается от предыдущей строки.
• Выполнить кумулятивную сумму __change_sides для обозначения каждого отдельного запуска
• Выполнить кумулятивную сумму значения для каждого отдельного запуска

Так что теперь у нас есть наша нативная функция Polars. Давайте проведем наши замеры еще раз.

К сожалению, мы не увидели улучшения в скорости выполнения нашей функции. Вероятно, это связано с количеством операторов over, которые нам приходится выполнять для вычисления значений последовательности. Однако, мы заметили ожидаемое уменьшение потребляемой памяти. Возможно, есть еще более эффективный способ реализации с помощью выражений Polars, но я пока не буду об этом беспокоиться.

Плагины Polars

Теперь давайте рассмотрим новые плагины Polars. Если вам нужен учебник по их настройке, обратитесь к документации. Здесь я больше сосредоточусь на конкретной реализации плагина. Сначала мы напишем наш алгоритм на Rust.

use polars::prelude::*;use pyo3_polars::derive::polars_expr;#[polars_expr(output_type=Float64)]fn calculate_runs(inputs: &[Series]) -> PolarsResult<Series> {    let values = inputs[0].f64()?;    let mut run_values: Vec<f64> = Vec::with_capacity(values.len());    let mut current_run_value = 0.0;    let mut run_is_positive = true;    for value in values {        match value {            None => {                run_values.push(current_run_value);            }            Some(value) => {                if value > 0.0 {                    if run_is_positive {                        current_run_value += value;                    } else {                        current_run_value = value;                        run_is_positive = true;                    }                } else if run_is_positive {                    current_run_value = value;                    run_is_positive = false;                } else {                    current_run_value += value;                }                run_values.push(current_run_value);            }        }    }    Ok(Series::from_vec("runs", run_values))}

Вы заметите, что это выглядит довольно похоже на алгоритм, который мы написали на Python. Здесь мы не используем никакой магии Rust! Мы обозначаем тип вывода с помощью макроса, предоставляемого Polars, и все. Затем мы можем зарегистрировать нашу новую функцию как выражение.

from polars import selectors as csfrom polars.utils.udfs import _get_shared_lib_locationlib = _get_shared_lib_location(__file__)@pl.api.register_expr_namespace("runs")class RunNamespace:    def __init__(self, expr: pl.Expr):        self._expr = expr    def calculate_runs(        self,    ) -> pl.Expr:        return self._expr.register_plugin(            lib=lib,            symbol="calculate_runs",            is_elementwise=False,            cast_to_supertypes=True,        )

А затем мы можем запустить его так:

from polars_extentsion import RunNamespacedf.select(  pl.col(feat_col).runs.calculate_runs().over(entity_col).alias("run_value")).collect()

Хорошо, теперь давайте проверим результаты!

Вот это уже по нашему вкусу! Мы получили улучшение скорости в 14 раз и снизили объем выделенной памяти с ~57MiB до ~8MiB.

Когда следует использовать плагины Polars

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

• Если вы легко можете написать очень быструю версию вашего расчета, используя собственные выражения Polars. Разработчики Polars очень умные. Я бы не ставил на то, что смогу написать функцию, работающую значительно быстрее, чем они. У Polars есть все необходимые инструменты. Используйте то, в чем они хороши!
• Если для вашего расчета нет естественной параллелизации. Например, если бы мы не выполняли указанную выше задачу для нескольких сущностей одновременно, наше ускорение было бы значительно меньше. Мы выиграли от скорости Rust и от естественной способности Polars применять нашу функцию Rust сразу к нескольким группам.
• Если вам не нужна высочайшая скорость работы или эффективность использования памяти. Многие согласятся, что написание на Rust гораздо сложнее и занимает больше времени, чем написание на Python. Поэтому, если вам все равно, занимает вашей функции 2 секунды или 200 мс, вам может не понадобиться использовать плагины.

Имея в виду вышесказанное, вот несколько требований, которые, на мой взгляд, подталкивают меня иногда использовать плагины:

• Скорость и использование памяти имеют большое значение. Недавно я переписал большую часть функционала конвейера данных в плагин Polars, потому что мы постоянно переключались между Polars и другими инструментами, и выделение памяти становилось слишком большим. Стало сложно запустить конвейер на желаемой инфраструктуре с желаемым объемом данных. Плагины позволили запустить тот же конвейер значительно быстрее и на более маленькой машине.
• У вас есть уникальное применение. Polars предоставляет множество встроенных функций. Но это общий инструментарий, который широко применим для многих задач. Иногда этот инструментарий не подходит для конкретной задачи, которую вы пытаетесь решить. В этом случае плагин может быть именно то, что вам нужно. Два наиболее частых примера, с которыми я сталкивался, это более интенсивные математические вычисления, такие как применение поперечной линейной регрессии, или последовательные (строчные) вычисления, как мы показали здесь.

Новая система плагинов идеально дополняет все вычисления на основе столбцов, которые Polars уже поддерживает из коробки. С этим дополнением Polars предоставляет красивую расширяемость возможностей. В дополнение к написанию собственных плагинов, ожидайте разработки интересных пакетов плагинов Polars, которые вы сможете использовать для расширения своих возможностей, не писав сами плагины!

Polars развивается быстро и создает волну. Ознакомьтесь с проектом, начните его использовать, обратите внимание на другие потрясающие функции, которые они будут выпускать, а может быть, в процессе выучите немного Rust!