带有文字类型的Python Polars 窗口函数

Question

假设我有一个带有 id 列的 DataFrame，如下所示：

df = pl.from_repr("""
┌─────┐
│ id  │
│ --- │
│ i64 │
╞═════╡
│ 1   │
│ 1   │
│ 1   │
│ 2   │
│ 2   │
│ 3   │
│ 3   │
└─────┘
""")

我想聚合 id 列上的运行计数，给出以下结果：

┌─────┬───────┐
│ id  ┆ count │
│ --- ┆ ---   │
│ i64 ┆ i64   │
╞═════╪═══════╡
│ 1   ┆ 1     │
│ 1   ┆ 2     │
│ 1   ┆ 3     │
│ 2   ┆ 1     │
│ 2   ┆ 2     │
│ 3   ┆ 1     │
│ 3   ┆ 2     │
└─────┴───────┘

我的尝试涉及创建一个虚拟列，我认为这产生了所需的结果，但似乎有点老套。

(
    df.with_columns(
        pl.lit(1).alias("ones")
    )
    .with_columns(
        (pl.col("ones").cum_sum().over("id")).alias("count")
    )
    .drop("ones")
)

但是当我尝试这个时：

(
    df.with_columns(
        (pl.lit(1).cum_sum().over("id")).alias("count")
    )
    .drop("ones")
)

我收到错误“ComputeError：窗口表达式的长度与组的长度不匹配”。

有更好的方法吗？我在上述尝试中缺少什么？

Answer 1

cumcount 似乎可以通过 1 班次调整来完成这项工作（当然，或者您想要的任何班次）。

还有arange：

df.with_columns(
   count = 1 + pl.col('id').cumcount().over('id'),
   count2 = pl.arange(1, 1 + pl.count()).over('id')
)