How to assign Exponential Moving Averages after groupby in python polars

Question

I have just started using polars in python and I'm coming from pandas. I would like to know how can I replicate the below pandas code in python polars

import pandas as pd
import polars as pl

df['exp_mov_avg_col'] = df.groupby('agg_col')['ewm_col'].transform(lambda x : x.ewm(span=14).mean())

I have tried the following:

df.groupby('agg_col').agg([pl.col('ewm_col').ewm_mean().alias('exp_mov_avg_col')])

but this gives me a list of exponential moving averages per provider, I want that list to be assigned to a column in original dataframe to the correct indexes, just like the pandas code does.

Answer 1

You can use window functions which apply an expression within a group defined by .over("group").

df = pl.DataFrame({
    "agg_col": [1, 1, 2, 3, 3, 3],
    "ewm_col": [1, 2, 3, 4, 5, 6]
})

(df.select([
    pl.all().exclude("ewm_col"),
    pl.col("ewm_col").ewm_mean(alpha=0.5).over("agg_col")
]))
 

Ouputs:

shape: (6, 2)
┌─────────┬──────────┐
│ agg_col ┆ ewm_col  │
│ ---     ┆ ---      │
│ i64     ┆ f64      │
╞═════════╪══════════╡
│ 1       ┆ 1.0      │
├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
│ 1       ┆ 1.666667 │
├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
│ 2       ┆ 3.0      │
├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
│ 3       ┆ 4.0      │
├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
│ 3       ┆ 4.666667 │
├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
│ 3       ┆ 5.428571 │
└─────────┴──────────┘