如何在Polars中计算当前行下方x行的rolling_sum？求更优方案

我需要实现基于当前行下方x行的滚动求和，而非常规的上方行。目前我有两种实现方案，想对比哪种更优（计算效率更高/代码更简洁可读）：

方案一：两次排序实现 #

先按分组键和排序键降序排列，计算滚动求和后再恢复原顺序：

import polars as pl

# 测试数据集
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"]
    })

(
    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")
    )
)

方案对比与优劣势分析 #

1. 计算效率 #

• Shift方案更高效：排序操作的时间复杂度为O(n log n)，两次排序会带来额外的CPU和内存开销；而Shift操作是线性时间O(n)，无需额外排序，在处理大型数据集时性能优势会非常明显。
• 排序方案：即使Polars的排序已经做了优化，仍无法避免两次排序的损耗，数据量越大，差距越显著。

2. 代码简洁性与可读性 #

• Shift方案更胜一筹：代码逻辑直接，无需理解“倒序计算再正序恢复”的绕弯逻辑，一眼就能看出是通过滚动求和加偏移来实现目标；排序方案的可读性相对较弱，需要额外梳理排序前后的行对应关系。

3. 注意事项 #

使用Shift方案时，要注意窗口大小和偏移量的对应关系：

• 如果需要计算当前行+下方n-1行的滚动求和（窗口大小为n），偏移量应为-(n-1)，比如窗口大小3对应shift(-2)；
• 如果需求是仅下方n行（不含当前行），则需要调整为pl.col("Close").shift(1).rolling_sum(n).over("Company")。

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