使用Polars为排序后的数据每行添加组内索引

初始实现 #

我需要为排序后的数据集每行添加组内索引，以下是最初的实现代码：

import polars as pl
import numpy as np

max_groups = 5
max_reps = 3

# 显示所有行方便查看
pl.Config.set_tbl_rows(max_groups * max_reps)

num_groups = np.random.randint(3, max_groups + 1)
unique_ids = np.random.randint(97, 123, num_groups)
repetitions = np.random.randint(1, max_reps + 1, num_groups)

id_col = "id"
data_col = "point"
index_col = "ixs"

# 生成数据
# 将整数转换为ASCII字符
ids = pl.Series(
    id_col,
    [c for n, id in zip(repetitions, unique_ids) for c in [chr(id)] * n],
)
data = pl.Series(
    data_col,
    np.random.rand(len(ids)),
)
df = pl.DataFrame([ids, data])

# 生成组内索引
df.sort(id_col, data_col).group_by(id_col).agg(
    pl.col(data_col), pl.int_range(pl.len()).alias(index_col)
).explode(data_col, index_col).sort(id_col, data_col)

执行后输出结果：

shape: (7, 3)
┌─────┬──────────┬─────┐
│ id  ┆ point    ┆ ixs │
│ --- ┆ ---      ┆ --- │
│ str ┆ f64      ┆ i64 │
╞═════╪══════════╪═════╡
│ g   ┆ 0.030686 ┆ 0   │
│ g   ┆ 0.322024 ┆ 1   │
│ k   ┆ 0.124792 ┆ 0   │
│ k   ┆ 0.289025 ┆ 1   │
│ s   ┆ 0.485742 ┆ 0   │
│ s   ┆ 0.689397 ┆ 1   │
│ u   ┆ 0.516705 ┆ 0   │
└─────┴──────────┴─────┘

优化方案 #

上述代码中进行了两次排序：分组前一次，分组后一次。通过在group_by中设置maintain_order=True，可以省去第二次排序，优化后的代码如下：

# 生成组内索引，利用maintain_order保留排序顺序
df.sort(id_col, data_col).group_by(id_col, maintain_order=True).agg(
    pl.col(data_col), pl.int_range(pl.len()).alias(index_col)
).explode(data_col, index_col)

性能说明 #

通过timeit简单测试发现，使用maintain_order=True通常比两次排序更高效，但性能优势并不显著。

内容的提问来源于stack exchange，提问作者bzm3r