如何在polars中创建条件增量列？

Question

我想在polars中创建一个有条件的递增列。
这个列应该从1开始，只有在满足某个条件（比如pl.col('code') == 'L'）时才递增。

import polars as pl
df = pl.DataFrame({'file': ['a.txt','a.txt','a.txt','a.txt','b.txt','b.txt','c.txt','c.txt','c.txt','c.txt','c.txt'],
                   'code': ['X','Y','Z','L','A','A','B','L','C','L','X']
                   })
df.with_columns(pl.int_range(start=1, end=pl.len()+1).over('file').alias('rrr')
                )

这段代码会产生一个简单的无条件递增列。但是我该如何添加条件呢？

Answer 1

看起来，pl.cum_sum()可以处理这个问题：

(df.with_columns(pl.when(pl.col('code') == 'L')
                  .then(1)
                  .otherwise(0)
                  .alias('xxx')
                 )
   .with_columns((pl.cum_sum('xxx')+1).over("file").alias('rrr')
                 )
 )

Answer 2

你也可以直接对布尔条件使用 cum_sum 函数：

df.with_columns(
   (pl.col('code') == 'L').cum_sum().over('file').alias('rrr') + 1
)

shape: (11, 3)
┌───────┬──────┬─────┐
│ file  ┆ code ┆ rrr │
│ ---   ┆ ---  ┆ --- │
│ str   ┆ str  ┆ u32 │
╞═══════╪══════╪═════╡
│ a.txt ┆ X    ┆ 1   │
│ a.txt ┆ Y    ┆ 1   │
│ a.txt ┆ Z    ┆ 1   │
│ a.txt ┆ L    ┆ 2   │
│ b.txt ┆ A    ┆ 1   │
│ b.txt ┆ A    ┆ 1   │
│ c.txt ┆ B    ┆ 1   │
│ c.txt ┆ L    ┆ 2   │
│ c.txt ┆ C    ┆ 2   │
│ c.txt ┆ L    ┆ 3   │
│ c.txt ┆ X    ┆ 3   │
└───────┴──────┴─────┘

Answer 3

我不太确定你具体期待什么样的输出，不过这里有一个例子，展示了如何仅在符合条件的行上增加计数，使用的是 cum_sum() 这个函数：

df.with_columns(
    pl.when(pl.col('code') == 'L').then(pl.lit(1)).otherwise(pl.lit(0)).alias('rrr')
).with_columns(
    pl.col('rrr').cum_sum().over('file') + 1
)

┌───────┬──────┬─────┐
│ file  ┆ code ┆ rrr │
│ ---   ┆ ---  ┆ --- │
│ str   ┆ str  ┆ i32 │
╞═══════╪══════╪═════╡
│ a.txt ┆ X    ┆ 1   │
│ a.txt ┆ Y    ┆ 1   │
│ a.txt ┆ Z    ┆ 1   │
│ a.txt ┆ L    ┆ 2   │
│ b.txt ┆ A    ┆ 1   │
│ b.txt ┆ A    ┆ 1   │
│ c.txt ┆ B    ┆ 1   │
│ c.txt ┆ L    ┆ 2   │
│ c.txt ┆ C    ┆ 2   │
│ c.txt ┆ L    ┆ 3   │
│ c.txt ┆ X    ┆ 3   │
└───────┴──────┴─────┘