如何在分组极坐标数据框中创建未聚合变量

Question

import polars as pl
import numpy as np

df_a = pl.DataFrame({'recipe': 'A','values': [1,2,3]})
df_b = pl.DataFrame({'recipe': 'B','values': [1,2,3]})

df = pl.concat([df_a, df_b], rechunk=True)

我想按组计算 cumsums，但是 .with_columns 不适用于分组数据...

df.group_by('recipe').with_columns(cum_val = np.cumsum(pl.col('values')))

在 R 中我可以执行此操作

library(tidyverse)
mtcars |> group_by(cyl) |> mutate(cum_wt=cumsum(wt))

Answer 1

在极坐标中，您可以直接使用窗口函数。尝试：

df.with_columns(
    pl.col("values")
    .cum_sum()
    .over("recipe")
)

结果是：

recipe,values
A,1
A,3
A,6
B,1
B,3
B,6