group_by 并在极坐标数据框中添加计数器列

Question

我有一个极坐标数据框

import polars as pl

df = pl.from_repr("""
┌─────────┬─────────────────────────────┐
│ item_id ┆ num_days_after_first_review │
│ ---     ┆ ---                         │
│ i64     ┆ i64                         │
╞═════════╪═════════════════════════════╡
│ 1       ┆ 3                           │
│ 1       ┆ 3                           │
│ 1       ┆ 10                          │
│ 2       ┆ 2                           │
│ 2       ┆ 2                           │
│ 3       ┆ 1                           │
│ 3       ┆ 5                           │
└─────────┴─────────────────────────────┘
""")

我想要有一列指示每个

item_id

相对于

num_days_after_first_review

;

的计数器

所以结果会是这样的

shape: (7, 3)
┌─────────┬─────────────────────────────┬─────┐
│ item_id ┆ num_days_after_first_review ┆ num │
│ ---     ┆ ---                         ┆ --- │
│ i64     ┆ i64                         ┆ i64 │
╞═════════╪═════════════════════════════╪═════╡
│ 1       ┆ 3                           ┆ 1   │
│ 1       ┆ 3                           ┆ 2   │
│ 1       ┆ 10                          ┆ 3   │
│ 2       ┆ 1                           ┆ 1   │
│ 2       ┆ 2                           ┆ 2   │
│ 3       ┆ 1                           ┆ 1   │
│ 3       ┆ 5                           ┆ 2   │
└─────────┴─────────────────────────────┴─────┘

Answer 1

一种方法是将

.over()

与

pl.count()

和

pl.arange()

一起使用

df.with_columns(
   pl.arange(1, pl.count() + 1)
     .over("item_id")
     .alias("num"))

shape: (7, 3)
┌─────────┬─────────────────────────────┬─────┐
│ item_id | num_days_after_first_review | num │
│ ---     | ---                         | --- │
│ i64     | i64                         | i64 │
╞═════════╪═════════════════════════════╪═════╡
│ 1       | 3                           | 1   │
│ 1       | 3                           | 2   │
│ 1       | 10                          | 3   │
│ 2       | 2                           | 1   │
│ 2       | 2                           | 2   │
│ 3       | 1                           | 1   │
│ 3       | 5                           | 2   │
└─────────┴─────────────────────────────┴─────┘