蟒极群后指数移动平均的赋值方法

Question

我刚开始在蟒蛇身上使用北极熊，我是从熊猫那里来的。我想知道如何在蟒蛇极地复制下面的熊猫代码

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