基于Polars处理分组时序数据生成LSTM输入序列的技术问询

针对你需要将带person分组的大型时序数据转换为LSTM适配短序列的需求，以下是两种比原生rolling更简洁高效的实现方案，直接满足你对(3,3,3,2)形状numpy数组或person键字典的要求：

方案一：Polars分组 + Numpy滑动窗口（推荐） #

利用Polars的分组能力结合Numpy的sliding_window_view，一步生成目标序列结构，代码简洁且性能优异（适合大数据集）。

示例代码 #

import polars as pl
import numpy as np
from numpy.lib.stride_tricks import sliding_window_view

# 构造示例数据（3个person，每人5条时序数据，2个特征）
df = pl.DataFrame({
    "person": [1,1,1,1,1, 2,2,2,2,2, 3,3,3,3,3],
    "feature1": np.arange(15),
    "feature2": np.arange(15, 30)
})

# 窗口参数定义
window_size = 3  # 序列步长
step_size = 1    # 滑动步长（默认1）

# 分组生成序列
grouped = df.group_by("person").agg(
    pl.col(["feature1", "feature2"]).map_batches(
        lambda x: sliding_window_view(x.to_numpy(), window_size, axis=0)
    ).alias("sequences")
)

# 转换为person为键的字典
seq_dict = {
    person: seqs 
    for person, seqs in zip(grouped["person"], grouped["sequences"])
}

# 转换为目标形状的numpy数组
seq_array = np.stack(list(seq_dict.values()))
print(seq_array.shape)  # 输出：(3, 3, 3, 2) 完全匹配需求

说明 #

• sliding_window_view直接对每个分组的特征矩阵生成滑动窗口，输出形状为(序列数, 窗口步长, 特征数)
• 分组后通过np.stack将所有person的序列合并，直接得到(人员数, 序列数, 窗口步长, 特征数)的目标数组
• 相比原生rolling，无需额外的窗口聚合后拆分操作，代码量减少一半以上

方案二：纯Polars实现（无额外依赖） #

如果希望完全基于Polars实现，可通过rolling结合列表聚合完成，但代码相对繁琐：

import polars as pl
import numpy as np

df = pl.DataFrame({
    "person": [1,1,1,1,1, 2,2,2,2,2, 3,3,3,3,3],
    "feature1": np.arange(15),
    "feature2": np.arange(15, 30)
})

window_size = 3

# 分组生成单特征滚动序列
grouped = df.group_by("person").agg(
    [
        pl.col(f).rolling_index(window_size).agg(pl.col(f).list()).alias(f"seq_{f}")
        for f in ["feature1", "feature2"]
    ]
).filter(pl.col("seq_feature1").list.lengths() > 0)

# 合并多特征为(窗口步长, 特征)的序列结构
grouped = grouped.with_columns(
    pl.col("seq_feature1")
    .list.eval(pl.element().zip_with(pl.col("seq_feature2").list.get(pl.int_range(0, pl.col("seq_feature1").list.lengths()))))
    .alias("sequences")
)

# 转换为字典和numpy数组的方式同方案一
seq_dict = {p: s.to_numpy() for p, s in zip(grouped["person"], grouped["sequences"])}
seq_array = np.stack(list(seq_dict.values()))

方案对比 #

内容的提问来源于stack exchange，提问作者Andrew P.