在 Python Polars 中对特定列使用 .over() 函数时如何对第一个元素进行累积求和

Question

我想知道是否有人可以启发我。

我正在尝试对特定列进行累积求和/分组（声明）。 我最初的想法是使用一些东西：

dataset.filter(pl.col("pty_nber").first().over("Declaration").cumsum() < 30 )

但不幸的是，它没有考虑 .over() ，只是对所有行进行累积求和。因此，它不是将 4 + 7 + 8 等相加，而是将其相加 4 + 4 + 4 + 4 + 7 ...

目标是至少显示一些完整的声明，而不是中间被删减。

提前致谢:)

示例请参见下面：

--> 并过滤掉超过特定阈值（例如 30）的 CUMSUM，以便确保没有一个声明不完整（即不包括该特定声明的所有 pty_nber）

Answer 1

数据

我简化了你的示例数据

import polars as pl
df = pl.DataFrame(
    {
        "declaration": [2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4],
        "item": [1, 2, 3, 1, 2, 3, 4, 1, 2, 3, 4],
        "pty_nber": [12, 12, 12, 9, 9, 9, 9, 16, 16, 16, 16],
    }
)

第一：

df.group_by("declaration", maintain_order=True).first().filter(
    pl.col("pty_nber").cum_sum() < 30
)


shape: (2, 3)
┌─────────────┬──────┬──────────┐
│ declaration ┆ item ┆ pty_nber │
│ ---         ┆ ---  ┆ ---      │
│ i64         ┆ i64  ┆ i64      │
╞═════════════╪══════╪══════════╡
│ 2           ┆ 1    ┆ 12       │
│ 3           ┆ 1    ┆ 9        │
└─────────────┴──────┴──────────┘

第二：

对于第二个，我使用 ITEM 列的帮助，因为我认为每个组的第一个元素等于第一个项目

df.filter(
    pl.when(pl.col("item") == 1).then(pl.col("pty_nber")).otherwise(0).cum_sum() < 30
)


shape: (7, 3)
┌─────────────┬──────┬──────────┐
│ declaration ┆ item ┆ pty_nber │
│ ---         ┆ ---  ┆ ---      │
│ i64         ┆ i64  ┆ i64      │
╞═════════════╪══════╪══════════╡
│ 2           ┆ 1    ┆ 12       │
│ 2           ┆ 2    ┆ 12       │
│ 2           ┆ 3    ┆ 12       │
│ 3           ┆ 1    ┆ 9        │
│ 3           ┆ 2    ┆ 9        │
│ 3           ┆ 3    ┆ 9        │
│ 3           ┆ 4    ┆ 9        │
└─────────────┴──────┴──────────┘