Python + Polars: добавление «индекса внутри группы» к каждой строке отсортированных данных

Вместо этого вы можете использовать .rank на pl.col(data_col) с .over(id_col).

rank отсчитывается от 1, поэтому вам придется вычесть 1.

df = df.sort(id_col, data_col).with_columns(
    ixs=(pl.col(data_col).rank("ordinal").over(id_col) - 1)
)

Вывод с помощью np.random.seed(0) (соответствует фрейму данных, который генерирует ваш код):

shape: (7, 3)
┌─────┬──────────┬─────┐
│ id  ┆ point    ┆ ixs │
│ --- ┆ ---      ┆ --- │
│ str ┆ f64      ┆ u32 │
╞═════╪══════════╪═════╡
│ a   ┆ 0.528895 ┆ 0   │
│ a   ┆ 0.568045 ┆ 1   │
│ a   ┆ 0.791725 ┆ 2   │
│ p   ┆ 0.437587 ┆ 0   │
│ p   ┆ 0.891773 ┆ 1   │
│ v   ┆ 0.383442 ┆ 0   │
│ v   ┆ 0.963663 ┆ 1   │
└─────┴──────────┴─────┘

Примечание. rank имеет множество режимов, из которых можно выбирать ничьи. Об этом вы можете прочитать по ссылке выше.

Вы уже были на правильном пути, используя pl.int_range . Однако конструкция group_by/agg здесь не нужна. Вместо этого можно использовать оконную функцию с pl.Expr.over. Это позволяет оценивать pl.int_range отдельно в каждой группе.

Это может выглядеть следующим образом.

(
    df
    .sort(id_col, data_col)
    .with_columns(
        pl.int_range(pl.len()).over(id_col)
    )
)

shape: (8, 3)
┌─────┬──────────┬─────────┐
│ id  ┆ point    ┆ literal │
│ --- ┆ ---      ┆ ---     │
│ str ┆ f64      ┆ i64     │
╞═════╪══════════╪═════════╡
│ m   ┆ 0.291593 ┆ 0       │
│ m   ┆ 0.60665  ┆ 1       │
│ q   ┆ 0.480906 ┆ 0       │
│ q   ┆ 0.545202 ┆ 1       │
│ q   ┆ 0.706958 ┆ 2       │
│ t   ┆ 0.156814 ┆ 0       │
│ y   ┆ 0.460135 ┆ 0       │
│ y   ┆ 0.631585 ┆ 1       │
└─────┴──────────┴─────────┘