Polars 计算时间序列中的并发事件数

Question

我有一个时间序列：

ts = pl.DataFrame(
    {
        "timestamp": [0, 1, 2, 3, 4, 5],
        "start_index": [0, None, None, 3, None, None],
        "end_index": [4, None, None, 4, None, None],
        }
    )

对于每个时间戳索引，我想计算并发事件的数量。例如，

• 在时间戳 0 处，我们有一个事件从时间戳 0 的 末尾开始，到时间戳 4 的末尾结束，跨度为 [1, 2, 3, 4]。
在时间戳 3 处，我们有另一个事件，从时间戳 3 的末尾开始，到 4 结束，sapn 为 [4]。

第一个事件的跨度在时间戳 3-4 期间与第二个事件重叠，因此时间戳 4 处的并发事件数为 2。我们将 0 分配给没有活动事件的观测值。

预期输出：

shape: (6, 4)
┌───────────┬─────────────┬───────────┬────────────────────────────┐
│ timestamp ┆ start_index ┆ end_index ┆ number_of_concurrent_event │
│ ---       ┆ ---         ┆ ---       ┆ ---                        │
│ i64       ┆ i64         ┆ i64       ┆ i64                        │
╞═══════════╪═════════════╪═══════════╪════════════════════════════╡
│ 0         ┆ 0           ┆ 4         ┆ 0                          │
│ 1         ┆ null        ┆ null      ┆ 1                          │
│ 2         ┆ null        ┆ null      ┆ 1                          │
│ 3         ┆ 3           ┆ 4         ┆ 1                          │
│ 4         ┆ null        ┆ null      ┆ 2                          │
│ 5         ┆ null        ┆ null      ┆ 0                          │
└───────────┴─────────────┴───────────┴────────────────────────────┘

Answer 1

我将开始创建一个新框架，将

start_index

 到 

end_index

 转换为单独的行。由此，我将与原始数据进行 

outer_coalesce

 连接，然后按时间戳进行分组，将原始时间戳的非空计数作为并发事件。

(
    ts
    .filter(pl.col('start_index').is_not_null() & pl.col('end_index').is_not_null())
    .select(
        original_timestamp="timestamp", 
        timestamp=pl.int_ranges(
            pl.col('start_index')+1,
            pl.col('end_index')+1
            )
        )
    .explode('timestamp')
    .join(ts, 
          on='timestamp', 
          how='outer_coalesce'
          )
    .group_by('timestamp',maintain_order=True)
    .agg(
        pl.col('start_index','end_index').drop_nulls().first(), 
        pl.col('original_timestamp').is_not_null().sum()
        )
)
shape: (6, 4)
┌───────────┬─────────────┬───────────┬────────────────────────────┐
│ timestamp ┆ start_index ┆ end_index ┆ number_of_concurrent_event │
│ ---       ┆ ---         ┆ ---       ┆ ---                        │
│ i64       ┆ i64         ┆ i64       ┆ u32                        │
╞═══════════╪═════════════╪═══════════╪════════════════════════════╡
│ 0         ┆ 0           ┆ 4         ┆ 0                          │
│ 1         ┆ null        ┆ null      ┆ 1                          │
│ 2         ┆ null        ┆ null      ┆ 1                          │
│ 3         ┆ 3           ┆ 4         ┆ 1                          │
│ 4         ┆ null        ┆ null      ┆ 2                          │
│ 5         ┆ null        ┆ null      ┆ 0                          │
└───────────┴─────────────┴───────────┴────────────────────────────┘