如何在Polars DataFrame的SELECT中访问前一行值？Heikin-Ashi蜡烛图计算

我有一个结构如下的DataFrame：

DF = { "open": [], "high":[], "low":[], "close":[] }

需要通过以下公式聚合这个DF：

# OPEN = (Open_n-1 + Close_n-1) / 2.0
# HIGH = Max(High_n, Open_n, Close_n)
# LOW = Min(Low_n, Open_n, Close_n)
# CLOSE = (Open_n + High_n + Low_n + Close_n) / 4.0

我在实现OPEN = (Open_n-1 + Close_n-1) / 2.0这一步遇到了问题，不知道如何访问前一行的值。考虑过逐行遍历或者用额外查询生成前一行占位符，但感觉这种方式不够优雅，希望能得到更合适的实现建议。

目前我已经实现的代码如下：

self.candle_data_ctx.price_candle_df.select([
    # OPEN = (Open_n-1 + Close_n-1) / 2.0
    #     TODO - implement
    # HIGH = Max(High_n, Open_n, Close_n)
    pl.max([pl.col('high_price'), pl.col('open_price'), pl.col('close_price')
                ]).alias('high_price'),
    # LOW = Min(Low_n, Open_n, Close_n)
    pl.min([pl.col('low_price'), pl.col('open_price'), pl.col('close_price')
                ]).alias('low_price'),
    # CLOSE = (Open_n + High_n + Low_n + Close_n) / 4.0
    (
       (
          pl.col('open_price') + pl.col('high_price') + pl.col('low_price') + pl.col('close_price')
        ) / 4.0
    ).alias('close_price'),
])

在Polars中，直接用shift()方法就能获取前一行的数据，完全不需要逐行遍历或者额外查询，代码会非常简洁优雅。

针对OPEN字段的实现，你可以这样写：

(pl.col('open_price').shift(1) + pl.col('close_price').shift(1)) / 2.0

shift(1)表示将当前列向下偏移1位，也就是取前一行的值。

下面是完整的实现代码：

self.candle_data_ctx.price_candle_df.select([
    # OPEN = (Open_n-1 + Close_n-1) / 2.0
    ((pl.col('open_price').shift(1) + pl.col('close_price').shift(1)) / 2.0).alias('open_price'),
    # HIGH = Max(High_n, Open_n, Close_n)
    pl.max([pl.col('high_price'), pl.col('open_price'), pl.col('close_price')]).alias('high_price'),
    # LOW = Min(Low_n, Open_n, Close_n)
    pl.min([pl.col('low_price'), pl.col('open_price'), pl.col('close_price')]).alias('low_price'),
    # CLOSE = (Open_n + High_n + Low_n + Close_n) / 4.0
    ((pl.col('open_price') + pl.col('high_price') + pl.col('low_price') + pl.col('close_price')) / 4.0).alias('close_price'),
])

需要注意的是，第一行的open_price会因为没有前一行数据而得到null，如果需要处理这个情况，可以用fill_null()方法来填充默认值，比如：

((pl.col('open_price').shift(1) + pl.col('close_price').shift(1)) / 2.0).fill_null(pl.col('open_price')).alias('open_price')

这样第一行会直接使用当前行的open_price作为默认值，你可以根据实际需求调整填充逻辑。

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