如何让Polars滚动窗口起始处未填满时返回Null而非缩短窗口？

问题说明 #

在使用Polars的滚动窗口计算时，默认行为是当窗口未完全填满时，会基于现有数据截断计算（如示例中2天窗口的第一天返回当天数值）。但需求是：直到窗口完全填满前，返回Null值。

示例代码 #

import polars as pl

dates = [
    "2020-01-01",
    "2020-01-02",
    "2020-01-03",
    "2020-01-04",
    "2020-01-05",
    "2020-01-06",
    "2020-01-01",
    "2020-01-02",
    "2020-01-03",
    "2020-01-04",
    "2020-01-05",
    "2020-01-06",
]
df = pl.DataFrame(
    {
        "dt": dates,
        "a": [3, 4, 2, 8, 10, 1, 1, 7, 5, 9, 2, 1],
        "b": ["Yes","Yes","Yes","Yes","Yes", "Yes", "No", "No", "No", "No", "No", "No"]
    }
).with_columns(
    pl.col("dt").str.strptime(pl.Date).set_sorted()
)
df = df.sort(by='dt')

# 原代码（默认行为）
original_result = df.rolling(
    index_column="dt", period="2d", group_by='b'
).agg(pl.col("a").mean().alias("ma_2d"))

当前默认结果 #

shape: (12, 3)
┌─────┬────────────┬───────┐
│ b   ┆ dt         ┆ ma_2d │
│ --- ┆ ---        ┆ ---   │
│ str ┆ date       ┆ f64   │
╞═════╪════════════╪═══════╡
│ Yes ┆ 2020-01-01 ┆ 3.0   │
│ Yes ┆ 2020-01-02 ┆ 3.5   │
│ Yes ┆ 2020-01-03 ┆ 3.0   │
│ Yes ┆ 2020-01-04 ┆ 5.0   │
│ Yes ┆ 2020-01-05 ┆ 9.0   │
│ Yes ┆ 2020-01-06 ┆ 5.5   │
│ No  ┆ 2020-01-01 ┆ 1.0   │
│ No  ┆ 2020-01-02 ┆ 4.0   │
│ No  ┆ 2020-01-03 ┆ 6.0   │
│ No  ┆ 2020-01-04 ┆ 7.0   │
│ No  ┆ 2020-01-05 ┆ 5.5   │
│ No  ┆ 2020-01-06 ┆ 1.5   │
└─────┴────────────┴───────┘

预期结果 #

希望窗口未填满时返回Null，例如：

• "Yes"组2020-01-01的ma_2d为Null
• "No"组2020-01-01的ma_2d为Null

解决方案 #

通过pl.when().then().otherwise()在聚合时判断窗口是否满足填满条件，这里以窗口内包含至少2个不同日期作为判断依据（对应2天窗口的填满要求）：

modified_result = df.rolling(
    index_column="dt", period="2d", group_by='b'
).agg(
    pl.when(pl.col("dt").n_unique() >= 2)
      .then(pl.col("a").mean())
      .otherwise(None)
      .alias("ma_2d")
)

print(modified_result)

修改后的结果 #

shape: (12, 3)
┌─────┬────────────┬───────┐
│ b   ┆ dt         ┆ ma_2d │
│ --- ┆ ---        ┆ ---   │
│ str ┆ date       ┆ f64   │
╞═════╪════════════╪═══════╡
│ Yes ┆ 2020-01-01 ┆ null  │
│ Yes ┆ 2020-01-02 ┆ 3.5   │
│ Yes ┆ 2020-01-03 ┆ 3.0   │
│ Yes ┆ 2020-01-04 ┆ 5.0   │
│ Yes ┆ 2020-01-05 ┆ 9.0   │
│ Yes ┆ 2020-01-06 ┆ 5.5   │
│ No  ┆ 2020-01-01 ┆ null  │
│ No  ┆ 2020-01-02 ┆ 4.0   │
│ No  ┆ 2020-01-03 ┆ 6.0   │
│ No  ┆ 2020-01-04 ┆ 7.0   │
│ No  ┆ 2020-01-05 ┆ 5.5   │
│ No  ┆ 2020-01-06 ┆ 1.5   │
└─────┴────────────┴───────┘

补充说明 #

如果需要更严格的时间跨度判断（比如确保窗口覆盖完整的2天时长，而非仅2个日期），可以调整判断条件为：

pl.when((pl.col("dt").max() - pl.col("dt").min()) >= pl.duration(days=2))
  .then(pl.col("a").mean())
  .otherwise(None)

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