Polars实现跨分组唯一ID：基于列值变化的cumsum需求

要生成跨所有分组的唯一ID列d，核心是先计算每个分组的前置偏移量（即当前分组之前所有分组的b最大值之和），再和组内的b相加。你之前的问题出在直接对组内最大值列c做cum_sum，会重复计算每个组内的多行最大值，导致偏移量错误。

实现步骤 #

1. 计算每个分组的b最大值，得到列c；
2. 提取唯一分组的c值，排序后计算累积和，再移位得到每个分组的前置偏移量（第一个分组偏移量为0）；
3. 将偏移量映射回原数据集，与b相加得到全局唯一IDd。

代码示例 #

假设你的输入数据如下：

import polars as pl

df = pl.DataFrame({
    "a": ["X", "X", "Y", "Y", "Y", "Z"],
    "b": [1, 2, 1, 2, 3, 1]
})

执行以下代码：

# 第一步：计算每个分组的b最大值
df = df.with_columns(c=pl.col("b").max().over("a"))

# 第二步：生成分组偏移量
group_offsets = (
    df.select("a", "c")
    .unique()  # 每个分组只保留一行最大值
    .sort("a")  # 按分组顺序排序（可根据需求调整排序字段）
    .with_columns(offset=pl.col("c").cum_sum().shift(1).fill_null(0))
)

# 第三步：合并偏移量并计算d列
df = df.join(group_offsets, on="a").with_columns(d=pl.col("b") + pl.col("offset"))

# 清理中间列（可选）
df = df.drop("c", "offset")

最终结果 #

关键说明 #

• 必须先提取唯一分组的最大值再做累积和，避免重复计算同一分组的最大值；
• shift(1)和fill_null(0)是为了让第一个分组的前置偏移量为0，后续分组的偏移量是前面所有分组的最大值总和；
• 分组排序的顺序会直接影响d列的生成结果，可根据实际业务需求调整sort("a")中的排序字段。

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