在Pandas中，我使用 .groupby().cumsum() 来生成一个计数列：

import pandas as pd

df = pd.DataFrame({'ID':['A','B','A','A','B','B','C','D','D','C']})

df['count'] = df['ID'].ne(df['ID'].shift()).groupby(df['ID']).cumsum()

  ID  count
0  A      1
1  B      1
2  A      2
3  A      2
4  B      2
5  B      2
6  C      1
7  D      1
8  D      1
9  C      2

在Polars中尝试同样的方法时，结果却全是1。这是为什么？

import polars as pl

df = pl.DataFrame({'ID':['A','B','A','A','B','B','C','D','D','C']})

df.with_columns(
    count = (pl.col('ID').ne(pl.col('ID').shift()).fill_null(True).cum_sum()).over('ID')
)

┌─────┬───────┐
│ ID  ┆ count │
│ --- ┆ ---   │
│ str ┆ u32   │
╞═════╪═══════╡
│ A   ┆ 1     │
│ B   ┆ 1     │
│ A   ┆ 1     │
│ A   ┆ 1     │
│ B   ┆ 1     │
│ B   ┆ 1     │
│ C   ┆ 1     │
│ D   ┆ 1     │
│ D   ┆ 1     │
│ C   ┆ 1     │
└─────┴───────┘

如何生成与Pandas方案相同的计数列？

解决方案

你用来生成“连续ID”的Pandas shift做法，也被称为游程编码（run-length encoding）。

Polars确实提供了一个专门的 .rle_id() 表达式来实现这个。

df.with_columns(count=pl.col.ID.rle_id()) 

# shape: (10, 2)
# ┌─────┬───────┐
# │ ID  ┆ count │
# │ --- ┆ ---   │
# │ str ┆ u32   │
# ╞═════╪═══════╡
# │ A   ┆ 0     │
# │ B   ┆ 1     │
# │ A   ┆ 2     │
# │ A   ┆ 2     │
# │ B   ┆ 3     │
# │ B   ┆ 3     │
# │ C   ┆ 4     │
# │ D   ┆ 5     │
# │ D   ┆ 5     │
# │ C   ┆ 6     │
# └─────┴───────┘

你接着可以对每个分组执行一个密集 .rank() 作为一个窗口函数来得到1..N的序列。

(df.with_columns(count=pl.col.ID.rle_id())
   .with_columns(pl.col("count").rank(method="dense").over("ID"))
)

shape: (10, 2)
┌─────┬───────┐
│ ID  ┆ count │
│ --- ┆ ---   │
│ str ┆ u32   │
╞═════╪═══════╡
│ A   ┆ 1     │
│ B   ┆ 1     │
│ A   ┆ 2     │
│ A   ┆ 2     │
│ B   ┆ 2     │
│ B   ┆ 2     │
│ C   ┆ 1     │
│ D   ┆ 1     │
│ D   ┆ 1     │
│ C   ┆ 2     │
└─────┴───────┘