Improving Polars performance

[armgilles]

Based sur un benchmark sur le projet

Problématique

Comment avoir les meilleurs performance en Polars pour la création de feature.

Solution

def fct_with_columns(data: pl.DataFrame) -> pl.DataFrame:
    """ """

    data = data.with_columns(pl.col("available_stands").shift(1).over("station_id").alias("available_stands_shift"))
    data = data.with_columns(pl.col("available_stands_shift").fill_null(pl.col("available_stands")))
    data = data.with_columns(transactions_out=(pl.col("available_stands") - pl.col("available_stands_shift")))

    data = data.with_columns(
        transactions_out=pl.when(pl.col("transactions_out") < 0).then(0).otherwise(pl.col("transactions_out"))
    )

    # Drop non useful column
    data = data.drop("available_stands_shift")

    return data


def fct_expr() -> pl.Expr:
    """ """
    available_stands_shift = (
        pl.col("available_stands").shift(1).over("station_id").fill_null(pl.col("available_stands"))
    )
    transactions_out = pl.col("available_stands") - available_stands_shift

    return pl.when(transactions_out < 0).then(0).otherwise(transactions_out).alias("transactions_out")

df = df.with_columns(fct_with_columns(df))
# and
df_lazy = fct_with_columns(df_lazy)

Les expressions ne semblent pas obligatoire pour optimiser le traitement, with_colomns fonctionne très bien avec un LazyFrame.

Dans ce cas, pour utiliser plusieurs fonctions on peut faire (si les fonctions ont en input et ouput un df) :

df = (
        df.pipe(fonction_1) 
        .pipe(fonction_2) 
        .pipe(fonction_3)
        .pipe(fonction_4)
    )