如何在Polars中基于事件表匹配并扩展时序DataFrame？

核心需求 #

将time_df中channel≠0的行按channel=0的分隔点分组，分别与events_df每行的event_table列表对齐：事件数量多于分组行数时截断，不足时填充None。

完整解决方案代码 #

import polars as pl

# 初始化示例数据
time_df = pl.DataFrame({
    "channel": [0, 1, 2, 1, 2, 0, 1],
    "time": [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7]
})

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

# 核心处理逻辑
result_df = (
    time_df
    # 1. 生成分组ID：channel=0时触发新分组，向前填充ID到后续行
    .with_columns(
        group_id=pl.when(pl.col("channel") == 0)
                  .then(pl.int_range(0, pl.count()))
                  .forward_fill()
    )
    # 2. 过滤channel=0的行，给每个分组内的行标记序号
    .filter(pl.col("channel") != 0)
    .with_row_index("row_idx", by="group_id")
    # 3. 展开事件列表并按分组+序号左连接，自动处理截断/补空
    .join(
        events_df
        .with_row_index("group_id")  # 给events_df分配和time_df对应的分组ID
        .explode("event_table")      # 展开事件列表
        .with_row_index("event_idx", by="group_id"),  # 标记事件的组内序号
        left_on=["group_id", "row_idx"],
        right_on=["group_id", "event_idx"],
        how="left"
    )
    # 4. 清理结果列
    .rename({"event_table": "event"})
    .drop("group_id", "row_idx", "event_idx", "time_right")
)

print(result_df)

步骤详解 #

1. 标记分组ID
   利用when+forward_fill给time_df中的非0行分配分组：channel=0的位置生成连续的分组编号，后续行继承这个编号，确保每个事件组的时序行归属正确。

2. 生成组内行号
   过滤掉channel=0的行后，用with_row_index(by="group_id")给每个分组内的行生成从0开始的递增序号，用于和事件列表的索引一一对应。

3. 对齐事件与时序行
   给events_df分配和time_df对应的分组ID，展开event_table后标记每个事件的组内序号，再通过左连接关联：

   • 事件数量多于分组行数：超出的事件因无对应时序行号，会被自动过滤（截断）
   • 事件数量少于分组行数：无对应事件的时序行，event列自动填充None

4. 清理结果
   重命名列名，删除中间生成的辅助列，得到目标结构。

测试事件数量不匹配的场景 #

将events_df替换为事件数量与分组行数不匹配的情况：

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

运行相同逻辑后，会得到期望的结果：

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

内容的提问来源于stack exchange，提问作者DJDuque