在Polars DataFrame中提取时序数据的所有可能子序列

问题背景 #

手里有大型时序数据集，每个独立时序由group_idx标识，time_idx列记录该时序存在的时间步（对应目标值存在）。最小示例代码：

import polars as pl

pl.Config(fmt_table_cell_list_len=10, fmt_str_lengths=80)

min_df = pl.DataFrame({
    "group_idx": [0, 1, 2, 3], "time_idx": [[0, 1, 2, 3], [2, 3], [0, 2, 3], [0,3]]
})

对应的DataFrame：

shape: (4, 2)
┌───────────┬──────────────┐
│ group_idx ┆ time_idx     │
│ ---       ┆ ---          │
│ i64       ┆ list[i64]    │
╞═══════════╪══════════════╡
│ 0         ┆ [0, 1, 2, 3] │
│ 1         ┆ [2, 3]       │
│ 2         ┆ [0, 2, 3]    │
│ 3         ┆ [0, 3]       │
└───────────┴──────────────┘

数据集时间范围共4个步长，但部分时序存在缺失时间步（如group_idx=3仅包含0和3）。需要从每个可能的时间步开始，按指定序列长度生成所有子序列，缺失时间步填充null。以sequence_length=3为例，预期输出：

result_df = pl.DataFrame({
    "group_idx": [0, 0, 1, 1, 2, 2, 3, 3], 
    "time_idx": [[0, 1, 2, 3], [0, 1, 2, 3], [2, 3], [2, 3], [0,2,3], [0,2,3], [0,3], [0,3]], 
    "sub_sequence": [[0,1,2], [1,2,3], [None, None, 2], [None, 2, 3], [0, None, 2], [None, 2, 3], [0, None, None], [None, None, 3]]
})

对应DataFrame：

shape: (8, 3)
┌───────────┬──────────────┬─────────────────┐
│ group_idx ┆ time_idx     ┆ sub_sequence    │
│ ---       ┆ ---          ┆ ---             │
│ i64       ┆ list[i64]    ┆ list[i64]       │
╞═══════════╪══════════════╪═════════════════╡
│ 0         ┆ [0, 1, 2, 3] ┆ [0, 1, 2]       │
│ 0         ┆ [0, 1, 2, 3] ┆ [1, 2, 3]       │
│ 1         ┆ [2, 3]       ┆ [null, null, 2] │
│ 1         ┆ [2, 3]       ┆ [null, 2, 3]    │
│ 2         ┆ [0, 2, 3]    ┆ [0, null, 2]    │
│ 2         ┆ [0, 2, 3]    ┆ [null, 2, 3]    │
│ 3         ┆ [0, 3]       ┆ [0, null, null] │
│ 3         ┆ [0, 3]       ┆ [null, null, 3] │
└───────────┴──────────────┴─────────────────┘

要求完全用Polars实现，适配超大数据集。

已尝试的失败方法 #

实际数据集explode()后约2亿行，group_idx和time_idx已排序，尝试以下代码时进程被终止：

(
 min_df.lazy()
     .with_columns(
         pl.col("time_idx").alias("time_idx_nulls")
     )
     .rolling(
          index_column='time_idx',
          group_by='group_idx',
          period=str(max_sequence_length) + 'i',
     )
     .agg(pl.col("time_idx_nulls"))
     .filter(pl.col('time_idx_nulls').list.len() == max_sequence_length)
)

高效解决方案 #

核心思路是：生成所有可能的起始位置，与每个group交叉连接，再为每个起始位置生成对应窗口的时间步，最后匹配原时序填充缺失值。全程采用Lazy模式，避免内存过载。

代码实现 #

import polars as pl

# 配置参数
sequence_length = 3
max_time_step = 3

# 生成所有合法的起始位置
starts = pl.DataFrame({"start": list(range(max_time_step - sequence_length + 2))})

# 核心逻辑
result = (
    min_df.lazy()
    # 交叉连接所有起始位置，每个group生成对应数量的子序列行
    .join(starts, how="cross")
    # 生成当前窗口的完整时间步列表
    .with_columns(
        window_time=pl.int_range(
            pl.col("start"), 
            pl.col("start") + sequence_length, 
            step=1, 
            dtype=pl.Int64
        )
    )
    # 将原time_idx转为Set提升匹配效率
    .with_columns(time_set=pl.col("time_idx").cast(pl.Set(pl.Int64)))
    # 匹配原时序值，缺失位置填充null
    .with_columns(
        sub_sequence=pl.col("window_time").list.eval(
            pl.when(pl.element().is_in(pl.col("time_set")))
            .then(pl.element())
            .otherwise(None)
        )
    )
    # 保留需要的列
    .select("group_idx", "time_idx", "sub_sequence")
    .collect()
)

print(result)

优化说明 #

1. 交叉连接生成起始位置：避免了对超大表的explode操作，内存占用可控
2. Set类型匹配：将time_idx转为pl.Set，大幅提升缺失值检查的效率，尤其适合长时序场景
3. Lazy模式执行：Polars会优化执行计划，分批次处理数据，适配亿级行的超大数据集

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