Polars 数据框：滚动和前瞻筛选

Question

我想计算一个叫做 rolling_sum 的东西，但不是计算当前行上面 x 行的数据，而是计算当前行下面 x 行的数据。

我的解决办法是先把数据表按照 descending=True 的方式排序，然后再用 rolling_sum，最后再把排序改回 descending=False。

我的解决方案：

import polars as pl

# Dummy dataset
df = pl.DataFrame({
        "Date": [1, 2, 3, 4, 5, 1, 2, 3, 4, 5],
        "Close": [-1, 1, 2, 3, 4, 4, 3, 2, 1, -1],
        "Company": ["A", "A", "A","A", "A",  "B", "B", "B", "B", "B"]
    })

# Solution using sort twice

(
    df
    .sort(by=["Company", "Date"], descending=[True, True])
    .with_columns(
        pl.col("Close").rolling_sum(3).over("Company").alias("Cumsum_lead")
    )
    .sort(by=["Company", "Date"], descending=[False, False])
)

有没有更好的解决办法呢？

我说的更好是指：

• 计算效率更高，或者
• 代码更少 / 更容易理解

谢谢！

编辑：

我刚想到另一个解决办法，可以完全避免排序或反转列：使用 shift。

(
    df
    .with_columns(
        pl.col("Close")
      .rolling_sum(3)
      .shift(-2)
      .over("Company").alias("Cumsum_lead"))
)

Answer 1

你可以不需要对行进行排序，而是可以通过两次反转特定的列来实现，使用的是 pl.Expr.reverse 这个功能。

(
    df
    .with_columns(
        pl.col("Close")
        .reverse().rolling_sum(3).reverse()
        .over("Company").alias("Cumsum_lead")
    )
)

为了让代码更易读，这个过程也可以放到一个辅助函数里。

def rolling_sum_lead(expr: pl.Expr, window_size: int) -> pl.Expr:
    return expr.reverse().rolling_sum(window_size).reverse()

(
    df
    .with_columns(
        rolling_sum_lead(pl.col("Close"), 3).over("Company").alias("Cumsum_lead")
    )
)

注意。 在我的电脑上，这个方法每次循环大约需要 124 微秒 ± 5.67 微秒，而使用 pl.DataFrame.sort 的方法则需要 205 微秒 ± 6.9 微秒。