根据条件增加值

根据当前的方法和预期的结果，我认为条件是cdcl_type等于"REC"或"GEC"。

预期输出可以通过以下方式获得。

• 对于满足条件的每个连续的行块，我们使用pl.Expr.rle_id条件表达式获取相应的 id。
• 我们使用 id 为每个这样的块创建一个增加的整数序列pl.int_range。
• 最后，我们将序列移位并加 1，并用 0 填充任何缺失的值。

df.with_columns(
    pl.when(
        pl.col("cdl_type").is_not_null()
    ).then(
        pl.int_range(
            pl.len()
        ).over(
            pl.col("cdl_type").is_in(["REC", "GEC"]).rle_id()
        )
    ).add(1).shift().fill_null(0)
)

shape: (9, 3)
┌────────┬──────────┬─────────┐
│ before ┆ cdl_type ┆ literal │
│ ---    ┆ ---      ┆ ---     │
│ 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       │
└────────┴──────────┴─────────┘

（免责声明：我之前写过这篇文章并发布它 - 刚刚意识到它与@Hericks 的回答有多么相似）。

由于您希望每次连续的“REC”或“GEC”组中断时重置计数，因此您需要将这个问题分成两部分。将逻辑重写为伪代码，您最终会得到：

• 前提：当当前行具有以下值之一"REC"或"GEC"...
• 然后：获取增量行数。
• 否则：增量计数应变为 0

一旦我们完成了这个逻辑，我们应该能够将所有内容向下移动一行来创建您想要的输出。

import polars as pl
from polars import col

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

print(
    df
    .with_columns(
        after=(
            pl.when(col('cdl_type').is_in(['GEC', 'REC']))
            .then(
                pl.int_range(pl.len()).add(1)
                .over(
                    col('cdl_type').is_in(['REC', 'GEC'])
                    .rle_id()
                )
            )
            .otherwise(0)
            .shift(1, fill_value=0)
        )
    )
)
# shape: (9, 3)
# ┌────────┬──────────┬───────┐
# │ before ┆ cdl_type ┆ after │
# │ ---    ┆ ---      ┆ ---   │
# │ i64    ┆ str      ┆ u32   │
# ╞════════╪══════════╪═══════╡
# │ 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 = pl.DataFrame({
    "before": [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    "cdl_type": ["REC", "REC", "GEC", None, None, None, "GEC", None, "REC", "GEC"],
})

^{（我最初的答案产生了所提供样本的预期输出，但对于长度 > 2 的 None 运行却不起作用）}

看起来您还想将第一个后续空值作为运行长度的一部分- 这可以通过 前向填充1 步来完成。

df.with_columns(
    pl.when(pl.col("cdl_type").is_not_null())
      .then(pl.col("cdl_type").is_not_null().rle_id())
      .forward_fill(limit=1)
      .rle_id()
      .alias("rle_id")
)

shape: (10, 3)
┌────────┬──────────┬────────┐
│ before ┆ cdl_type ┆ rle_id │
│ ---    ┆ ---      ┆ ---    │
│ i64    ┆ str      ┆ u32    │
╞════════╪══════════╪════════╡
│ 0      ┆ REC      ┆ 0      │
│ 0      ┆ REC      ┆ 0      │
│ 0      ┆ GEC      ┆ 0      │
│ 0      ┆ null     ┆ 0      │ # <-
│ 0      ┆ null     ┆ 1      │
│ 0      ┆ null     ┆ 1      │
│ 0      ┆ GEC      ┆ 2      │
│ 0      ┆ null     ┆ 2      │ # <-
│ 0      ┆ REC      ┆ 3      │
│ 0      ┆ GEC      ┆ 3      │
└────────┴──────────┴────────┘

rle()给出一个包含每个 {len,value} 的结构

df.select(
    pl.when(pl.col("cdl_type").is_not_null())
      .then(pl.col("cdl_type").is_not_null().rle_id())
      .forward_fill(limit=1)
      .rle()
)

shape: (4, 1)
┌───────────┐
│ cdl_type  │
│ ---       │
│ struct[2] │
╞═══════════╡
│ {4,0}     │
│ {2,null}  │
│ {2,2}     │
│ {2,4}     │
└───────────┘

这些len值被赋予int_ranges()并被展平以创建计数列。

df.with_columns(
    pl.int_ranges(
        pl.when(pl.col("cdl_type").is_not_null())
          .then(pl.col("cdl_type").is_not_null().rle_id())
          .forward_fill(limit=1)
          .rle()
          .struct.field("len")
    )
    .flatten()
    .alias("after")
)

shape: (10, 3)
┌────────┬──────────┬───────┐
│ before ┆ cdl_type ┆ after │
│ ---    ┆ ---      ┆ ---   │
│ i64    ┆ str      ┆ i64   │
╞════════╪══════════╪═══════╡
│ 0      ┆ REC      ┆ 0     │
│ 0      ┆ REC      ┆ 1     │
│ 0      ┆ GEC      ┆ 2     │
│ 0      ┆ null     ┆ 3     │
│ 0      ┆ null     ┆ 0     │
│ 0      ┆ null     ┆ 1     │ # <- NOT OK
│ 0      ┆ GEC      ┆ 0     │
│ 0      ┆ null     ┆ 1     │
│ 0      ┆ REC      ┆ 0     │
│ 0      ┆ GEC      ┆ 1     │
└────────┴──────────┴───────┘

然后我们将剩余的空值设置回 0

df.with_columns(
    pl.int_ranges(
        pl.when(pl.col("cdl_type").is_not_null())
          .then(pl.col("cdl_type").is_not_null().rle_id())
          .forward_fill(limit=1)
          .rle()
          .struct.field("len")
    )
    .flatten()
    .alias("after")
).with_columns(
    pl.when(pl.col("cdl_type").shift().is_not_null())
      .then(pl.col("after"))
      .otherwise(0)
)

shape: (10, 3)
┌────────┬──────────┬───────┐
│ before ┆ cdl_type ┆ after │
│ ---    ┆ ---      ┆ ---   │
│ i64    ┆ str      ┆ i64   │
╞════════╪══════════╪═══════╡
│ 0      ┆ REC      ┆ 0     │
│ 0      ┆ REC      ┆ 1     │
│ 0      ┆ GEC      ┆ 2     │
│ 0      ┆ null     ┆ 3     │
│ 0      ┆ null     ┆ 0     │
│ 0      ┆ null     ┆ 0     │ # <- OK
│ 0      ┆ GEC      ┆ 0     │
│ 0      ┆ null     ┆ 1     │
│ 0      ┆ REC      ┆ 0     │
│ 0      ┆ GEC      ┆ 1     │
└────────┴──────────┴───────┘

采用这种方法的原因是它避免使用.over()- 并且如果处理更大的数据框会产生更快的结果。

• https://github.com/pola-rs/polars/issues/19089

df.with_columns(
    after = pl.int_range(pl.len()).over(
        pl.col.cdl_type
        .is_in(["REC", "GEC"])
        .fill_null(pl.int_range(pl.len()) + 2)
        .rle_id()
    )
    .add(~pl.col.cdl_type.is_null()).shift().fill_null(0)
)

shape: (9, 3)
┌────────┬──────────┬───────┐
│ 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     │
└────────┴──────────┴───────┘