groupby transform lambda ewm using polars

Question

I'm trying to find faster alternative for the following code using polars

df['ewm'] = df.groupby(['outlet', 'product'])['sales'].transform(lambda x: x.shift(shift).ewm(com=10).mean())

Please find MWE example below

df = pd.DataFrame({"outlet": [1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3, 3], 
                   "product": [1, 1, 1, 2, 2, 2, 3, 3, 3, 1, 1, 1, 2, 2, 2, 3, 3, 3, 1, 1, 1, 2, 2, 2, 3, 3, 3], 
                   "sales": [1, 2, 3, 4, 5, 6, 7, 8, 9, 1, 2, 3, 4, 5, 6, 7, 8, 9, 1, 2, 3, 4, 5, 6, 7, 8, 9]})

outlet  product sales   ewm
0   1   1   1   NaN
1   1   1   2   1.00
2   1   1   3   1.52
3   1   2   4   NaN
4   1   2   5   4.00
5   1   2   6   4.52
6   1   3   7   NaN
7   1   3   8   7.00
8   1   3   9   7.52
....

Can you add a small MWE example showing your input data and your expected result? It can be a pandas example. — ritchie46, Sep 27 '22 at 17:09

score 2 · Answer 1 · answered Sep 28 '22 at 01:54

Here's the equivalent in Polars. (I'll also include the result from Pandas for comparison.)

shift = 1
df['ewm_pandas'] = df.groupby(['outlet', 'product'])['sales'].transform(lambda x: x.shift(shift).ewm(com=10).mean())

df = pl.from_pandas(df)
(
    df
    .with_column(
        pl.col('sales')
        .shift(shift)
        .ewm_mean(com=10)
        .over(['outlet', 'product'])
        .alias('ewn_polars')
    )
)

shape: (27, 5)
┌────────┬─────────┬───────┬────────────┬────────────┐
│ outlet ┆ product ┆ sales ┆ ewm_pandas ┆ ewn_polars │
│ ---    ┆ ---     ┆ ---   ┆ ---        ┆ ---        │
│ i64    ┆ i64     ┆ i64   ┆ f64        ┆ f64        │
╞════════╪═════════╪═══════╪════════════╪════════════╡
│ 1      ┆ 1       ┆ 1     ┆ null       ┆ null       │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 1      ┆ 1       ┆ 2     ┆ 1.0        ┆ 1.0        │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 1      ┆ 1       ┆ 3     ┆ 1.52381    ┆ 1.52381    │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 1      ┆ 2       ┆ 4     ┆ null       ┆ null       │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 1      ┆ 2       ┆ 5     ┆ 4.0        ┆ 4.0        │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 1      ┆ 2       ┆ 6     ┆ 4.52381    ┆ 4.52381    │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 1      ┆ 3       ┆ 7     ┆ null       ┆ null       │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 1      ┆ 3       ┆ 8     ┆ 7.0        ┆ 7.0        │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 1      ┆ 3       ┆ 9     ┆ 7.52381    ┆ 7.52381    │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 2      ┆ 1       ┆ 1     ┆ null       ┆ null       │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 2      ┆ 1       ┆ 2     ┆ 1.0        ┆ 1.0        │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 2      ┆ 1       ┆ 3     ┆ 1.52381    ┆ 1.52381    │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 2      ┆ 2       ┆ 4     ┆ null       ┆ null       │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 2      ┆ 2       ┆ 5     ┆ 4.0        ┆ 4.0        │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 2      ┆ 2       ┆ 6     ┆ 4.52381    ┆ 4.52381    │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 2      ┆ 3       ┆ 7     ┆ null       ┆ null       │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 2      ┆ 3       ┆ 8     ┆ 7.0        ┆ 7.0        │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 2      ┆ 3       ┆ 9     ┆ 7.52381    ┆ 7.52381    │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 3      ┆ 1       ┆ 1     ┆ null       ┆ null       │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 3      ┆ 1       ┆ 2     ┆ 1.0        ┆ 1.0        │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 3      ┆ 1       ┆ 3     ┆ 1.52381    ┆ 1.52381    │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 3      ┆ 2       ┆ 4     ┆ null       ┆ null       │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 3      ┆ 2       ┆ 5     ┆ 4.0        ┆ 4.0        │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 3      ┆ 2       ┆ 6     ┆ 4.52381    ┆ 4.52381    │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 3      ┆ 3       ┆ 7     ┆ null       ┆ null       │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 3      ┆ 3       ┆ 8     ┆ 7.0        ┆ 7.0        │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 3      ┆ 3       ┆ 9     ┆ 7.52381    ┆ 7.52381    │
└────────┴─────────┴───────┴────────────┴────────────┘

groupby transform lambda ewm using polars

1 Answers1