如何以Polars原生方式根据Buy_Signal生成多列并聚合Returns？

核心解法 #

利用Polars的向量化操作结合列表推导，一次性生成所有需要的聚合列，无需手动指定列名或循环拼接DataFrame：

import polars as pl
import numpy as np

# 构造示例DataFrame
df = pl.DataFrame({
    "Buy_Signal": [1, 0, 1, 0, 0], 
    "Returns": np.random.normal(0, 0.1, 5),
})

# 获取所有Buy_Signal为1的行索引
buy_positions = df["Buy_Signal"].arg_true().to_list()

# 生成每个信号对应的区间聚合列（以求和为例，可替换为mean/std等）
result_df = df.with_columns(
    [
        pl.col("Returns").slice(pos).sum().alias(f"Port_{idx+1}")
        for idx, pos in enumerate(buy_positions)
    ]
)

print(result_df)

关键细节说明 #

• 定位信号位置：arg_true()是Polars原生的高效方法，直接返回所有布尔值为True的行索引，比手动遍历效率更高
• 区间聚合：slice(pos)会从索引pos的位置开始，截取到DataFrame末尾的所有Returns值，后续直接调用聚合函数（sum()/mean()等）完成计算，全程为向量化操作
• 动态列生成：通过列表推导式自动根据信号数量生成对应的列名（如Port_1、Port_2），无需手动硬编码，适配任意数量的Buy Signal场景

替代方案（窗口函数实现） #

如果偏好窗口函数风格，也可以用滚动聚合实现相同效果：

result_df = df.with_columns(
    [
        pl.col("Returns")
          .rolling_sum(window_size=len(df)-pos, offset=-pos)
          .alias(f"Port_{idx+1}")
        for idx, pos in enumerate(buy_positions)
    ]
)

两种方法性能接近，slice的写法更直观易懂，适合大多数场景。

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