寻求Polars中基于Buy_Signal的重叠分组高效可读实现方案

针对你从Pandas迁移到Polars，需要基于Buy_Signal列中值为1的位置生成重叠分组并展开Returns的需求，这里提供一种更高效且可读性更强的方案，避免了原方案中的join操作，更适合大规模数据集：

实现代码 #

import polars as pl

# 原始数据
df = pl.DataFrame({
    "Buy_Signal": [1, 0, 1, 0, 1, 0, 0], 
    "Returns": [0.01, 0.02, 0.03, 0.02, 0.01, 0.00, -0.01],
})

# 1. 获取所有Buy_Signal为1的行的索引
buy_indices = df.select(pl.int_range(0, pl.len()).filter(pl.col("Buy_Signal") == 1)).to_series().to_list()

# 2. 生成每个组对应的Returns序列，再展开成目标格式
result = (
    pl.DataFrame({
        "group": range(1, len(buy_indices) + 1),
        "Returns": [df["Returns"][idx:] for idx in buy_indices]
    })
    .explode("Returns")
)

print(result)

方案优势 #

1. 性能更优：直接利用Polars的切片操作提取每个组的Returns序列，避免了原方案中join带来的额外开销，在大规模数据集上性能提升明显。
2. 可读性强：代码逻辑清晰直观，从获取信号位置到生成分组序列再到展开，每一步意图明确，易于维护。
3. 原生Polars操作：核心操作均为Polars原生方法，充分利用其向量化计算的优势，适配大数据场景。

输出验证 #

运行上述代码后，输出结果与目标完全一致：

shape: (15, 2)
┌───────┬─────────┐
│ group ┆ Returns │
│ ---   ┆ ---     │
│ i64   ┆ f64     │
╞═══════╪═════════╡
│ 1     ┆ 0.01    │
│ 1     ┆ 0.02    │
│ 1     ┆ 0.03    │
│ 1     ┆ 0.02    │
│ 1     ┆ 0.01    │
│ 1     ┆ 0.0     │
│ 1     ┆ -0.01   │
│ 2     ┆ 0.03    │
│ 2     ┆ 0.02    │
│ 2     ┆ 0.01    │
│ 2     ┆ 0.0     │
│ 2     ┆ -0.01   │
│ 3     ┆ 0.01    │
│ 3     ┆ 0.0     │
│ 3     ┆ -0.01   │
└───────┴─────────┘

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