Polars技术问题：如何通过谓词查找前序极值对应的价格？

要实现需求，可以利用Polars的窗口函数结合条件判断，精准获取当前行之前最近同类型（'top'/'bottom'）spikes对应的price值，仅在当前行spikes为'top'或'bottom'时填充，其余行保留None。

完整代码 #

import polars as pl

# 初始化原DataFrame
df = pl.DataFrame({
    'price': [1, 2, 4, 2, 1, 2, 3, 2],
    'spikes': [None, None, 'top', None, 'bottom', None, 'top', None]
})

# 新增目标列
df = df.with_columns(
    pl.when(pl.col('spikes').is_in(['top', 'bottom']))
      .then(pl.col('price').shift().over(pl.col('spikes')))
      .alias('prev_spikes_prices')
)

print(df)

代码说明 #

• pl.col('spikes').is_in(['top', 'bottom'])：筛选出spikes为'top'或'bottom'的行，仅对这些行进行值填充。
• pl.col('price').shift().over(pl.col('spikes'))：以spikes为分组依据，在每个分组内对price列执行向前移位操作，获取当前行之前最近同类型spikes对应的price值。
• 未满足条件的行自动填充None，完全匹配期望输出。

输出结果 #

shape: (8, 3)
┌───────┬────────┬──────────────────┐
│ price ┆ spikes ┆ prev_spikes_prices │
│ ---   ┆ ---    ┆ ---              │
│ i64   ┆ str    ┆ i64              │
╞═══════╪════════╪══════════════════╡
│ 1     ┆ null   ┆ null             │
│ 2     ┆ null   ┆ null             │
│ 4     ┆ top    ┆ null             │
│ 2     ┆ null   ┆ null             │
│ 1     ┆ bottom ┆ null             │
│ 2     ┆ null   ┆ null             │
│ 3     ┆ top    ┆ 4                │
│ 2     ┆ null   ┆ null             │
└───────┴────────┴──────────────────┘

内容的提问来源于stack exchange，提问作者Joris Medeišis