Utilizzare i plugin di Polars per un incremento della velocità del 14x con Rust

Incrementa la velocità del tuo codice del 14x utilizzando i plugin di Polars con Rust

Raggiungere alta velocità al di fuori della libreria nativa Polars

Introduzione

Polars sta conquistando il mondo grazie alla sua velocità, efficienza di memoria e meravigliosa API. Se vuoi sapere quanto è potente, non devi guardare oltre i benchmark di DuckDB. E questi non utilizzano nemmeno la versione più recente di Polars.

Tuttavia, nonostante tutte le cose straordinarie che Polars può fare, tradizionalmente non è stata una soluzione migliore di Pandas per effettuare TUTTI i calcoli che potresti voler fare. Ci sono alcune eccezioni in cui Polars non ha ottenuto prestazioni migliori. Con il recente rilascio del sistema di plugin di Polars per Rust, però, potrebbe non essere più così.

Plugin di Polars

Cosa significa esattamente un plugin di Polars? È semplicemente un modo per creare le tue espressioni Polars utilizzando nativamente Rust ed esporle come espressioni utilizzando un namespace personalizzato. Ti consente di sfruttare la velocità di Rust e applicarla al tuo Polars DataFrame per eseguire calcoli in modo da sfruttare la velocità e gli strumenti integrati forniti da Polars.

Diamo un’occhiata ad alcuni esempi concreti.

• AI Coding Google Bard è un buon sviluppatore Python?
• Privacy dei dati e il suo impatto sulla gestione
• Mettere a prova il nostro stage coltivare la prossima generazione nel mezzo della marcia aziendale dell’IA

Calcoli sequenziali

Un’area in cui Polars sembra mancare di alcune funzionalità sono le operazioni che richiedono la conoscenza del valore precedente di un DataFrame. I calcoli di natura sequenziale non sono sempre facili o efficienti da scrivere con le espressioni native di Polars. Diamo un’occhiata a un esempio specifico.

Abbiamo l’algoritmo seguente per calcolare il valore cumulativo di un array di numeri per una determinata sequenza, definita come un insieme di numeri che hanno lo stesso segno. Ad esempio:

┌───────┬───────────┐│ value ┆ run_value ││ ---   ┆ ---       ││ i64   ┆ i64       │╞═══════╪═══════════╡│ 1     ┆ 1         │   # La prima sequenza inizia qui│ 2     ┆ 3         ││ 3     ┆ 6         ││ -1    ┆ -1        │   # La sequenza si resetta qui│ -2    ┆ -3        ││ 1     ┆ 1         │   # La sequenza si resetta qui└───────┴───────────┘

Quindi vogliamo avere la somma cumulativa di una colonna che si resetta ogni volta che il segno del valore cambia da positivo a negativo o da negativo a positivo.

Iniziamo con una versione di base scritta in 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)

Iteriamo su una serie, calcolando il valore attuale della sequenza in ogni posizione e restituendo una nuova serie di Pandas.

Benchmarking

Prima di procedere, creeremo alcuni benchmark. Misureremo sia la velocità di esecuzione che il consumo di memoria utilizzando pytest-benchmark e pytest-memray. Configureremo il problema in modo da avere una colonna di entità, una colonna di tempo e una colonna di caratteristiche. L’obiettivo è calcolare i valori della sequenza per ogni entità dei dati nel tempo. Imposteremo il numero di entità e il numero di timestamp a 1.000, ottenendo così un DataFrame con 1.000.000 di righe.

Quando eseguiamo la nostra implementazione di Pandas rispetto al nostro benchmark utilizzando la funzionalità groupby apply di Pandas, otteniamo i seguenti risultati:

Implementazione Naive di Polars

Ora abbiamo il nostro benchmark. Vediamo come implementare questa stessa funzionalità in Polars. Inizieremo con una versione molto simile che verrà applicata mappando la funzione su un oggetto Polars GroupBy.

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

Ora vediamo come questa soluzione si confronta con il nostro benchmark originale di Pandas.

Bene, questo non ha funzionato molto bene. Non dovrebbe sorprendere però. Gli autori di Polars hanno reso molto chiaro che l’approccio comune di groupby apply in Pandas non è un modo efficiente per eseguire calcoli in Polars. Qui si nota. Sia la velocità che il consumo di memoria sono peggiori rispetto alla nostra implementazione originale di Pandas.

Implementazione Polars Expression

Scriviamo adesso questa stessa funzione come espressioni native di Polars. Questo è il modo preferito e ottimizzato per lavorare con Polars. L’algoritmo avrà un aspetto leggermente diverso. Ma ecco cosa ho ideato per calcolare lo stesso output.

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("__"))    )

Una breve spiegazione di cosa stiamo facendo qui:

• Troviamo tutte le righe in cui la caratteristica è positiva
• Troviamo tutte le righe in cui la colonna __is_positive è diversa dalla riga precedente.
• Prendiamo una somma cumulativa di __change_sides per contrassegnare ogni run distinta
• Prendiamo una somma cumulativa del valore su ogni run distinta

Ora abbiamo la nostra funzione nativa di Polars. Facciamo di nuovo il nostro benchmark.

Sfortunatamente non abbiamo visto un miglioramento nella velocità di esecuzione della nostra funzione. Ciò è probabilmente dovuto al numero di dichiarazioni over che dobbiamo fare per calcolare i valori di esecuzione. Tuttavia, abbiamo riscontrato una riduzione della memoria attesa. Potrebbe esserci un modo ancora migliore per implementare questo con le espressioni Polars, ma per ora non mi preoccuperò.

Plugin Polars

Ora diamo un’occhiata ai nuovi plugin Polars. Se desideri un tutorial su come configurarli, consulta la documentazione qui. Qui mostrerò principalmente un’implementazione specifica di un plugin. Per prima cosa scriveremo il nostro algoritmo in 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))}

<p.noterai abbiamo="" alcuna="" all'algoritmo="" assomiglia="" basta.="" che="" come="" da="" di="" e="" espressione.

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

E poi possiamo eseguirlo così:

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

Okay, ora diamo un’occhiata ai risultati!

Ora è più come piace a me! Abbiamo ottenuto un miglioramento della velocità del 14x e siamo passati da ~57MiB a ~8MiB di memoria allocata.

Quando utilizzare i plugin di Polars

Ora che ho dimostrato il potere dell’utilizzo dei plugin, parliamo di quando non dovresti usarli. Ecco alcuni motivi per cui potrei non utilizzare i plugin (ciascuno con i suoi avvertimenti):

• Se puoi facilmente scrivere una versione molto veloce del tuo calcolo utilizzando le espressioni native di Polars. Gli sviluppatori di Polars sono molto intelligenti. Non scommetterei denaro sul fatto che io possa scrivere una funzione significativamente più veloce di loro. Gli strumenti per Polars sono lì. Approfitta di ciò in cui sono bravi!
• Se non esiste una parallelizzazione naturale per il tuo calcolo. Ad esempio, se non stessimo eseguendo il problema sopra su più entità, il miglioramento della velocità sarebbe probabilmente stato significativamente inferiore. Abbiamo beneficiato sia dalla velocità di Rust che dalla capacità naturale di Polars di applicare la nostra funzione Rust su più gruppi contemporaneamente.
• Se non hai bisogno di prestazioni di primo livello in termini di velocità o memoria. Molti concordano sul fatto che scrivere in Rust sia molto più difficile e richieda più tempo rispetto a scrivere in Python. Quindi, se non ti importa se la tua funzione richiede 2 secondi per essere eseguita invece di 200 ms, potresti non aver bisogno di utilizzare i plugin.

Tenendo presente le considerazioni sopra riportate, ecco alcuni requisiti che mi spingono ad utilizzare i plugin alcune volte:

• La velocità e la memoria sono molto importanti. Ho recentemente riscritto gran parte delle funzionalità di un flusso di dati in un plugin di Polars perché stavamo alternando tra Polars e altri strumenti e le allocazioni di memoria stavano diventando troppo grandi. Stava diventando difficile eseguire il flusso di lavoro sull’infrastruttura desiderata con la quantità di dati che volevamo. I plugin hanno reso facile eseguire lo stesso flusso di lavoro in molto meno tempo e su una macchina molto più piccola.
• Hai un caso d’uso unico. Polars fornisce così tante funzioni incorporate. Ma è un set di strumenti generico che è ampiamente applicabile a molti problemi. A volte, questo set di strumenti non è specificamente applicabile al problema che stai cercando di risolvere. In questo caso, un plugin potrebbe essere esattamente ciò di cui hai bisogno. Due degli esempi più comuni che ho incontrato sono calcoli matematici più intensi, come l’applicazione di una regressione lineare trasversale, o calcoli sequenziali (basati su righe) come mostrato qui.

Il nuovo sistema di plugin è il complemento perfetto a tutti i calcoli basati sulle colonne che Polars già supporta. Con questa aggiunta, Polars consente un’estensibilità splendida alle sue capacità. Oltre a scrivere i tuoi plugin, fai attenzione a alcuni pacchetti di plugin Polars in sviluppo che puoi utilizzare per estendere le tue capacità senza dover scrivere plugin tu stesso!

Polars si sta muovendo rapidamente e sta facendo parlare di sé. Dai un’occhiata al progetto, inizia a usarlo, tieni d’occhio quali altre straordinarie funzionalità pubblicheranno e magari inizia a imparare un po’ di Rust mentre ci sei!