如何在滚动窗口中检查所有值是否为正并生成标记列

需求：给定包含一列数值的Polars DataFrame，设置滚动窗口长度为2，新增一列is_positive：当窗口内所有元素均为正数时标记为1，否则为0；窗口元素不足2个时（即首行）标记为null。

实现方案 #

方法1：向量化滚动窗口最小值判断（推荐，效率更高）

利用滚动窗口的最小值特性：如果窗口内所有元素都是正数，那么窗口最小值必然大于0。通过这个逻辑可以高效实现需求，无需遍历每个窗口：

import polars as pl

# 创建示例DataFrame
df = pl.DataFrame({
    "col": [1, -1, 1, -1, 1, 1, -1, 1]
})

# 新增标记列
df = df.with_columns(
    is_positive=pl.col("col")
    .rolling(window_size=2, min_periods=2)  # 窗口长度2，至少需要2个元素才计算
    .min()  # 取窗口最小值
    .gt(0)  # 判断最小值是否大于0，返回布尔值
    .cast(pl.Int64)  # 布尔值转成整数1/0
)

print(df)

方法2：窗口遍历判断（直观易懂）

通过apply遍历每个滚动窗口，直接判断窗口内所有元素是否为正：

import polars as pl

df = pl.DataFrame({
    "col": [1, -1, 1, -1, 1, 1, -1, 1]
})

df = df.with_columns(
    is_positive=pl.col("col")
    .rolling(window_size=2, min_periods=2)
    .apply(lambda window: 1 if (window > 0).all() else 0)
    .cast(pl.Int64)
)

print(df)

输出结果 #

两种方法都会得到符合需求的结果：

┌─────┬─────────────┐
│ col ┆ is_positive │
│ --- ┆ ---         │
│ i64 ┆ i64         │
╞═════╪═════════════╡
│ 1   ┆ null        │
│ -1  ┆ 0           │
│ 1   ┆ 0           │
│ -1  ┆ 0           │
│ 1   ┆ 0           │
│ 1   ┆ 1           │
│ -1  ┆ 0           │
│ 1   ┆ 0           │
└─────┴─────────────┘

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