使用极坐标进行分组变换 lambda ewm

Question

我正在尝试使用极坐标为以下代码找到更快的替代方案

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

请在下面找到 MWE 示例

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
....

Answer 1

这是 Polars 中的等效项。 （我还将包括 Pandas 的结果以进行比较。）

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