python-polars - 如何在 python polars 中的 groupby 之后分配指数移动平均线

Question

我刚刚开始在 python 中使用极坐标，我来自熊猫。我想知道如何在 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())

我尝试了以下方法：

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

但这给了我每个提供者的指数移动平均值列表，我希望将该列表分配给原始数据框中的列以正确的索引，就像熊猫代码一样。

Answer 1

您可以使用窗口函数在由 定义的组内应用表达式.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")
]))
 

输出：

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 │
└─────────┴──────────┘