基于条件递增Polars DataFrame列值并统计条件满足次数

需求说明 #

根据cdl_type列的内容生成after列，规则如下：

• 连续的REC/GEC序列：从0开始依次递增
• 非null序列后的第一个null：值为序列最后一个值+1
• 连续null中除第一个外的其他null：值为0
• 新的REC/GEC序列重新从0开始递增
• 单个REC/GEC后的null：值为1

示例数据 #

import polars as pl

df = pl.DataFrame({
    "before": [0, 0, 0, 0, 0, 0, 0, 0, 0],
    "cdl_type": ["REC", "REC", "GEC", None, None, "GEC", None, "REC", "GEC"],
})

期望输出 #

┌────────┬──────────┬───────┐
│ before ┆ cdl_type ┆ after │
│ ---    ┆ ---      ┆ ---   │
│ i64    ┆ str      ┆ i64   │
╞════════╪══════════╪═══════╡
│ 0      ┆ REC      ┆ 0     │
│ 0      ┆ REC      ┆ 1     │
│ 0      ┆ GEC      ┆ 2     │
│ 0      ┆ null     ┆ 3     │
│ 0      ┆ null     ┆ 0     │
│ 0      ┆ GEC      ┆ 0     │
│ 0      ┆ null     ┆ 1     │
│ 0      ┆ REC      ┆ 0     │
│ 0      ┆ GEC      ┆ 1     │
└────────┴──────────┴───────┘

正确实现方案 #

核心思路是通过标记序列分组、区分特殊null的方式，精准计算对应值：

df = df.with_columns(
    # 标记非null序列组：连续的REC/GEC归为同一组，null不参与分组
    group=pl.when(pl.col("cdl_type").is_in(["REC", "GEC"]))
            .then(pl.col("cdl_type").is_null().cumsum())
            .otherwise(None),
    # 标记是否是非null序列后的第一个null
    is_first_null=pl.col("cdl_type").is_null() & pl.col("cdl_type").shift(1).is_in(["REC", "GEC"])
).with_columns(
    # 非null序列内的递增计数
    seq_count=pl.when(pl.col("cdl_type").is_in(["REC", "GEC"]))
                .then(pl.int_range(0, pl.count()).over("group"))
                .otherwise(None),
    # 获取每个非null序列的最后一个计数值，用于第一个null的计算
    last_seq_count=pl.col("seq_count").last().over("group")
).with_columns(
    after=pl.when(pl.col("cdl_type").is_in(["REC", "GEC"]))
            .then(pl.col("seq_count"))
            .when(pl.col("is_first_null"))
            .then(pl.col("last_seq_count") + 1)
            .otherwise(0)
).drop("group", "is_first_null", "seq_count", "last_seq_count")

print(df)

代码解释 #

1. 分组标记：用cdl_type.is_null().cumsum()生成非null序列的分组ID，连续的非null会被分到同一组，null则不会计入分组。
2. 第一个null标记：通过shift(1)判断当前null的前一个值是否为非null，精准标记序列结束后的第一个null。
3. 序列内计数：在每个非null分组内，用int_range生成从0开始的递增序列，实现序列内的数值递增。
4. 生成after列：
   • 非null值直接取序列内的递增计数
   • 第一个null取对应分组的最后计数+1
   • 其他情况（非第一个null）统一设为0

原有方案的问题 #

原代码用rle()处理非null序列的计数，但未正确处理序列后的null逻辑，无法区分第一个null和后续连续null，导致输出不符合预期。

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