将数据框拆分为另一个数据框的范围

Question

我有两个数据框，里面的数据看起来像这样：

import polars as pl

data = {"channel": [0, 1, 2, 1, 2, 0, 1], "time": [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7]}
time_df = pl.DataFrame(data)

data = {
    "time": [10.0, 10.5],
    "event_table": [["start_1", "stop_1", "start_2", "stop_2"], ["start_3"]],
}
events_df = pl.DataFrame(data)

在这个数据中，time_df里的channel 0表示一个新的“事件表”开始。我想从event_df中的channel 0开始，把每一行的event_table展开，最终得到像这样的结果：

data = {
    "channel": [1, 2, 1, 2, 1],
    "time": [0.2, 0.3, 0.4, 0.5, 0.7],
    "event": ["start_1", "stop_1", "start_2", "stop_2", "start_3"],
}
result_df = pl.DataFrame(data)

我现在的做法是先把第一个数据框中的所有channel 0去掉，然后展开第二个数据框，最后用hstack把两个数据框合并在一起。如果我的数据没有问题，这样做是可以的。

但实际上，事件表可能会有更多（或更少）的事件。在这种情况下，我想要“截断”展开的结果（或者用空值填充），比如说：

import polars as pl

data = {"channel": [0, 1, 2, 1, 2, 0, 1], "time": [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7]}
time_df = pl.DataFrame(data)

data = {
    "time": [10.0, 10.5],
    "event_table": [["start_1", "stop_1", "start_2"], ["start_3", "stop_3"]],
}
events_df = pl.DataFrame(data)

data = {
    "channel": [1, 2, 1, 2, 1],
    "time": [0.2, 0.3, 0.4, 0.5, 0.7],
    "event": ["start_1", "stop_1", "start_2", None, "start_3"],
}
result_df = pl.DataFrame(data)

我非常感谢任何帮助。

Answer 1

我不确定自己是不是想得太复杂了，但：

最初想到的办法是给事件编号，这样它们就能和频道的行索引对齐。

shape: (2, 2)
┌──────────────────────────────────┬───────┐
│ event_table                      ┆ index │
│ ---                              ┆ ---   │
│ list[str]                        ┆ u32   │
╞══════════════════════════════════╪═══════╡ 
│ ["start_1", "stop_1", "start_2"] ┆ 0     │ # [1, 2, 3]
│ ["start_3", "stop_3"]            ┆ 5     │ # [6, 7]
└──────────────────────────────────┴───────┘

1. 运行长度编码

我们找到每个频道0的起始索引和运行长度：

rle_df = (
   time_df
    .with_row_index("start")
    .filter(channel = 0)
    .select(
       "start",
       len = pl.col("start").diff().shift(-1) - 1 
    )
    .with_row_index()
)

shape: (2, 3)
┌───────┬───────┬──────┐
│ index ┆ start ┆ len  │
│ ---   ┆ ---   ┆ ---  │
│ u32   ┆ u32   ┆ i64  │
╞═══════╪═══════╪══════╡
│ 0     ┆ 0     ┆ 4    │
│ 1     ┆ 5     ┆ null │
└───────┴───────┴──────┘

如果 len < event_len，我们想把这个作为新的长度。

2. 截断

我们把运行长度和 events_df 连接起来，然后使用

• .min_horizontal() 来获取 "最小" 长度
• 用 .list.head() 进行截断

events_with_index = (
   events_df
    .with_row_index()
    .join(rle_df, on="index", how="left")
    .select(
       pl.col("event_table").list.head(
          pl.min_horizontal(pl.col("len"), pl.col("event_table").list.len())
       ),
       index = pl.col("start")
    )
    .explode("event_table")
    .with_columns(
       pl.col("index") + pl.col("index").cum_count().over("index")
    )
)

在 .explode() 之后，我们得到了：

┌─────────────┬───────┐
│ event_table ┆ index │
│ ---         ┆ ---   │
│ str         ┆ u32   │
╞═════════════╪═══════╡
│ start_1     ┆ 0     │
│ stop_1      ┆ 0     │
│ start_2     ┆ 0     │
│ start_3     ┆ 5     │
│ stop_3      ┆ 5     │
└─────────────┴───────┘

然后加上 .cum_count()，我们得到了：

shape: (5, 2)
┌─────────────┬───────┐
│ event_table ┆ index │
│ ---         ┆ ---   │
│ str         ┆ u32   │
╞═════════════╪═══════╡
│ start_1     ┆ 1     │ # index of 1st channel 0 (plus 1)
│ stop_1      ┆ 2     │
│ start_2     ┆ 3     │
│ start_3     ┆ 6     │ # index of 2nd channel 0 (plus 1)
│ stop_3      ┆ 7     │
└─────────────┴───────┘

3. 连接

当索引对齐后，我们可以进行 join 操作，并过滤掉频道0的行。

• 左连接会让我们得到 "用空值填充" 的效果。

(time_df
  .with_row_index()
  .join(
     events_with_index,
     on = "index",
     how = "left"
  )     
  .filter(pl.col("channel") != 0)
)

shape: (5, 4)
┌───────┬─────────┬──────┬─────────────┐
│ index ┆ channel ┆ time ┆ event_table │
│ ---   ┆ ---     ┆ ---  ┆ ---         │
│ u32   ┆ i64     ┆ f64  ┆ str         │
╞═══════╪═════════╪══════╪═════════════╡
│ 1     ┆ 1       ┆ 0.2  ┆ start_1     │
│ 2     ┆ 2       ┆ 0.3  ┆ stop_1      │
│ 3     ┆ 1       ┆ 0.4  ┆ start_2     │
│ 4     ┆ 2       ┆ 0.5  ┆ null        │
│ 6     ┆ 1       ┆ 0.7  ┆ start_3     │
└───────┴─────────┴──────┴─────────────┘