如何在极坐标中的组上添加范围列

Question

我有：

df = pl.DataFrame({'key':['a','a','a','b','b','b'],'a':[2,4,6,1,2,3]})
print(df)
shape: (6, 2)
┌─────┬─────┐
│ key ┆ a   │
│ --- ┆ --- │
│ str ┆ i64 │
╞═════╪═════╡
│ a   ┆ 2   │
│ a   ┆ 4   │
│ a   ┆ 6   │
│ b   ┆ 1   │
│ b   ┆ 2   │
│ b   ┆ 3   │
└─────┴─────┘

我想在 key 定义的每个组中添加一个具有递增整数的列

df = pl.DataFrame({'key':['a','a','a','b','b','b'],'a':[2,4,6,1,2,3], 'r': [1,2,3,1,2,3]})
print(df)
shape: (6, 3)
┌─────┬─────┬─────┐
│ key ┆ a   ┆ r   │
│ --- ┆ --- ┆ --- │
│ str ┆ i64 ┆ i64 │
╞═════╪═════╪═════╡
│ a   ┆ 2   ┆ 1   │
│ a   ┆ 4   ┆ 2   │
│ a   ┆ 6   ┆ 3   │
│ b   ┆ 1   ┆ 1   │
│ b   ┆ 2   ┆ 2   │
│ b   ┆ 3   ┆ 3   │
└─────┴─────┴─────┘

我该怎么做？

Answer 1

在 github 问题上找到了与 Polars 相关的答案：使用

cum_count

df.with_columns(pl.col('c').cum_count().over("key").alias("r"))