选择 Polars 中具有连续日期的第一行

Question

我有一个 Polars DataFrame，其中包含每月日期和值列。我想只选择月份连续的第一行，删除其余的行。例如，在表中，这将是前 5 行。我尝试使用 .group_by_dynamic 但没有成功。有什么想法如何做到这一点？谢谢！

from datetime import datetime
df = pl.DataFrame(
    {
        "time": pl.datetime_range(
            start=datetime(2020, 12, 1),
            end=datetime(2023, 12, 1),
            interval="1mo",
            eager=True,
        ),
        "n": range(37),
    }
)

df=df.sample(n=10, seed=0)  
df=df.sort(["time"], descending=False)
df

时间：日期时间[μs]	n
2021-01-01 00:00:00	1
2021-02-01 00:00:00	2
2021-03-01 00:00:00	3
2021-04-01 00:00:00	4
2021-05-01 00:00:00	5
2022-03-01 00:00:00	15
2022-07-01 00:00:00	19
2023-04-01 00:00:00	28
2023-05-01 00:00:00	29
2023-06-01 00:00:00	30
2023-11-01 00:00:00	35

Answer 1

一种可能性是使用来自

filter

和

Expr.diff

的布尔表达式工匠来

cum_min

:

df.filter(pl.col('n').diff().fill_null(1).eq(1).cast(int).cum_min().cast(bool))

输出：

shape: (5, 2)
┌─────────────────────┬─────┐
│ time                ┆ n   │
│ ---                 ┆ --- │
│ datetime[μs]        ┆ i64 │
╞═════════════════════╪═════╡
│ 2021-01-01 00:00:00 ┆ 1   │
│ 2021-02-01 00:00:00 ┆ 2   │
│ 2021-03-01 00:00:00 ┆ 3   │
│ 2021-04-01 00:00:00 ┆ 4   │
│ 2021-05-01 00:00:00 ┆ 5   │
└─────────────────────┴─────┘

Answer 2

“连续”月份的要求对我来说并不完全清楚 - 我们是否完全忽略日期时间的

day

和

year

部分而只比较月份？

因此，如果稍微简化一下示例，以下是获取第一个连续行的方法。我将其基于第

n

列的第一个连续值，以便您可以看到该算法。您使用

Expr.shift

获取下一个值进行比较，然后使用

Expr.cum_sum

过滤掉行：

df.filter(
    (pl.col('n').shift(-1) - pl.col('n')).cum_sum() == 1
)

|time:datetime[μs]   |   n|
|--------------------|----|
|2021-01-01 00:00:00 |   1|
|2021-02-01 00:00:00 |   2|
|2021-03-01 00:00:00 |   3|
|2021-04-01 00:00:00 |   4|
|2021-05-01 00:00:00 |   5|

现在，您可以轻松地将其更改为比较月份

Answer 3

您可以使用

.rle_id()

来识别/标记“条纹”。

如果它们是连续的，则月份差异为 1。 （除了序列中的第一行，因此后退填充（1））

diff = pl.col("time").dt.month().diff()

df.with_columns(
   rle_id = pl.when(diff == 1).then(diff.rle_id()).backward_fill(1)
)

shape: (11, 3)
┌─────────────────────┬─────┬────────┐
│ time                ┆ n   ┆ rle_id │
│ ---                 ┆ --- ┆ ---    │
│ datetime[μs]        ┆ i64 ┆ u32    │
╞═════════════════════╪═════╪════════╡
│ 2021-01-01 00:00:00 ┆ 1   ┆ 1      │
│ 2021-02-01 00:00:00 ┆ 2   ┆ 1      │
│ 2021-03-01 00:00:00 ┆ 3   ┆ 1      │
│ 2021-04-01 00:00:00 ┆ 4   ┆ 1      │
│ 2021-05-01 00:00:00 ┆ 5   ┆ 1      │
│ 2022-03-01 00:00:00 ┆ 15  ┆ null   │
│ 2022-07-01 00:00:00 ┆ 19  ┆ null   │
│ 2023-04-01 00:00:00 ┆ 28  ┆ 5      │
│ 2023-05-01 00:00:00 ┆ 29  ┆ 5      │
│ 2023-06-01 00:00:00 ┆ 30  ┆ 5      │
│ 2023-11-01 00:00:00 ┆ 35  ┆ null   │
└─────────────────────┴─────┴────────┘

在本例中，第一个 id 恰好是

1

，因为我们在第一行中有一个匹配项。

我们可以使用

.rank()

来给我们保证从 1 开始的连续数字。

diff = pl.col("time").dt.month().diff()
rle_id = pl.when(diff == 1).then(diff.rle_id()).backward_fill(1)

df.filter(rle_id.rank("dense") == 1)

shape: (5, 2)
┌─────────────────────┬─────┐
│ time                ┆ n   │
│ ---                 ┆ --- │
│ datetime[μs]        ┆ i64 │
╞═════════════════════╪═════╡
│ 2021-01-01 00:00:00 ┆ 1   │
│ 2021-02-01 00:00:00 ┆ 2   │
│ 2021-03-01 00:00:00 ┆ 3   │
│ 2021-04-01 00:00:00 ┆ 4   │
│ 2021-05-01 00:00:00 ┆ 5   │
└─────────────────────┴─────┘