group_by and add a counter column in polars dataframe

Question

I have a polars dataframe

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

I would like to have a column that indicates a counter for each item_id with respect to num_days_after_first_review;

so the result will be like

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

One approach is to use pl.int_range(pl.len()) with .over()

df.with_columns(
   pl.int_range(1, pl.len() + 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   │
└─────────┴─────────────────────────────┴─────┘