Polars聚合中获取指定位置元素及map_groups优化问询

问题背景 #

在Polars中使用groupby、rolling、group_by_dynamic等聚合操作时，需要获取聚合后数据的指定位置元素（比如峰谷检测场景中，滚动窗口的中心位置元素）。当前通过map_groups实现但运行较慢，存在两个核心问题：

1. 能否通过一次map_groups调用返回多列，避免多次调用以提升速度？
2. 如何安全访问精确索引位置，而非仅依赖first()/last()？尝试col.Date.list[30]、col.Date.get(1)等写法均报错。

问题1：单次map_groups返回多列 #

可以让map_groups的自定义函数返回pl.DataFrame，一次性生成多列结果，减少Python函数调用开销，提升性能。

改造后的示例代码：

def process_window(date_series, ema_series):
    middle_loc = date_series.len() // 2
    peak_date = date_series[middle_loc]
    is_peak = ema_series[middle_loc] == ema_series.max()
    return pl.DataFrame({
        "Peak_Date": [peak_date],
        "Peak": [is_peak]
    })

peak_df = ema_df.group_by_dynamic(col.Date, every=minimum_bar_size, period="14mo").agg(
    Last_Date=col.Date.last(),
    SPX=col.SPX.last(),
    EMA_6mo=col.EMA_6mo.last(),
    Test=col.Date.len()//2,
    # 单次map_groups返回多列结果
    pl.map_groups([col.Date, col.EMA_6mo], process_window).alias("window_results")
).unnest("window_results")  # 展开嵌套的DataFrame列
.drop("Date")
.rename({"Last_Date":"Date"})
.unique(subset="Date")
.sort(by="Date")

问题2：安全访问精确索引位置 #

Polars聚合后的单列是原类型的Series而非List，直接用list[]或get()会报错。正确做法是先转成List，再结合长度判断避免越界：

纯表达式实现（推荐，性能远高于map_groups）

# 安全获取窗口中间位置的日期，处理长度不足的边界情况
peak_date_expr = pl.when(col.Date.len() >= 1)
                    .then(col.Date.to_list().get(col.Date.len() // 2))
                    .otherwise(None)

# 安全判断是否为峰值
is_peak_expr = pl.when(col.Date.len() >= 1)
                  .then(col.EMA_6mo.to_list().get(col.Date.len() // 2) == col.EMA_6mo.max())
                  .otherwise(False)

peak_df = ema_df.group_by_dynamic(col.Date, every=minimum_bar_size, period="14mo").agg(
    Last_Date=col.Date.last(),
    SPX=col.SPX.last(),
    EMA_6mo=col.EMA_6mo.last(),
    Test=col.Date.len()//2,
    Peak_Date=peak_date_expr,
    Peak=is_peak_expr
).drop("Date")
.rename({"Last_Date":"Date"})
.unique(subset="Date")
.sort(by="Date")

报错原因说明

• col.Date.list[30]：聚合后的col.Date不是List列，必须先通过to_list()转为List类型才能用索引访问。
• col.Date.get(1)：get()方法针对单个元素的列，而非聚合后的序列，因此会触发越界错误。

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