在Polars分组DataFrame中计算距上次登录行数（前置行设为None）

给定如下Polars DataFrame：

import polars as pl

df = pl.DataFrame({
    "user_id": [1, 1, 1, 2, 2, 2], 
    "login": [False, True, False, False, False, True]
})

需要添加一列count，实现：

• 计算每行距离用户上次登录的行数
• 用户首次登录前的行设为None

期望输出如下：

┌─────────┬───────┬───────┐
│ user_id ┆ login ┆ count │
│ ---     ┆ ---   ┆ ---   │
│ i64     ┆ bool  ┆ i64   │
╞═════════╪═══════╪═══════╡
│ 1       ┆ false ┆ null  │
│ 1       ┆ true  ┆ 0     │
│ 1       ┆ false ┆ 1     │
│ 2       ┆ false ┆ null  │
│ 2       ┆ false ┆ null  │
│ 2       ┆ true  ┆ 0     │
└─────────┴───────┴───────┘

尝试过参考单组场景的实现，但无法适配按user_id分组的情况。

可以通过分组窗口函数结合累计计数、分组标记来实现，以下是两种可行方案：

方案一：基于登录行基准计算 #

通过标记登录行的位置，向前填充最近登录基准后计算差值：

import polars as pl

df = pl.DataFrame({
    "user_id": [1, 1, 1, 2, 2, 2], 
    "login": [False, True, False, False, False, True]
})

result = df.with_columns(
    # 分组内生成每行的序号
    pl.int_range(0, pl.count()).over("user_id").alias("row_num"),
    # 仅在登录行记录当前序号，其余为null
    pl.when(pl.col("login")).then(pl.int_range(0, pl.count())).over("user_id").alias("login_row")
).with_columns(
    # 向前填充最近的登录行序号
    pl.col("login_row").fill_null(strategy="forward").over("user_id"),
    # 获取分组内第一个登录行的序号
    pl.col("login_row").first().over("user_id").alias("first_login")
).with_columns(
    # 首次登录前的行设为null，其余计算当前行与最近登录行的序号差
    pl.when(pl.col("row_num") >= pl.col("first_login"))
      .then(pl.col("row_num") - pl.col("login_row"))
      .alias("count")
).drop("row_num", "login_row", "first_login")

print(result)

方案二：基于登录区间分组 #

通过登录行为的累计计数划分区间，在每个区间内计算相对行数：

import polars as pl

df = pl.DataFrame({
    "user_id": [1, 1, 1, 2, 2, 2], 
    "login": [False, True, False, False, False, True]
})

result = df.with_columns(
    # 按用户分组，统计累计登录次数，形成登录区间标记
    pl.col("login").cumsum().over("user_id").alias("login_grp")
).with_columns(
    # 首次登录前（login_grp=0）设为null，其余区间内从0开始计数
    pl.when(pl.col("login_grp") == 0)
      .then(None)
      .otherwise(pl.int_range(0, pl.count()).over(["user_id", "login_grp"]))
      .alias("count")
).drop("login_grp")

print(result)

两种方案都能实现需求，方案二更简洁，核心是利用login_grp区分首次登录前后的区间，再在每个区间内计算相对行数。

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