Python Polars 滚动计数

Question

Polars中有一些已知的滚动函数，即rolling_mean()、rolling_apply()和rolling_max()。但是，如果我想统计每个窗口中某个值出现的次数，该怎么做？

假设我们现在有一个 LazyFrame:

df = pl.LazyFrame({"Date": ["2023-01-01", "2023-01-02", "2023-01-03", "2023-01-05", "2023-01-10", "2023-01-11", "2023-01-12"], "Pattern": [True, True, False, True, False, False, True]})
┌────────────┬─────────┐
│ Date       ┆ Pattern │
│ ---        ┆ ---     │
│ str        ┆ bool    │
╞════════════╪═════════╡
│ 2023-01-01 ┆ true    │
│ 2023-01-02 ┆ true    │
│ 2023-01-03 ┆ false   │
│ 2023-01-05 ┆ true    │
│ 2023-01-10 ┆ false   │
│ 2023-01-11 ┆ false   │
│ 2023-01-12 ┆ true    │
└────────────┴─────────┘

对于n = 3且pattern = True，期望的结果是:

┌────────────┬─────────┐
│ Date       ┆ Count   │
│ ---        ┆ ---     │
│ str        ┆ int     │
╞════════════╪═════════╡
│ 2023-01-01 ┆ null    │
│ 2023-01-02 ┆ null    │
│ 2023-01-03 ┆ 2       │
│ 2023-01-05 ┆ 2       │
│ 2023-01-10 ┆ 1       │
│ 2023-01-11 ┆ 1       │
│ 2023-01-12 ┆ 1       │
└────────────┴─────────┘

我尝试在 Pattern 列上使用rolling_sum()，但由于我的列是 bool 类型，使用此类函数只会产生错误。

在 pandas 中，这可以通过以下方式实现:

df["Pattern"].apply(lambda x: x == True).rolling(3, min_periods = 0).sum()

使用极坐标的正确方法是什么？并且，如何将解决方案推广到 bool 值 True 值以外的列，例如 False 和分类数据？

Answer 1

I have tried using rolling_sum() over the column Pattern, yet since my column is of Boolean type, using such functions would only yield an error.

您可以使用rolling_sum，只需首先使用Expr.cast将 bool 表达式转换为数字类型即可(False 值转换为 0，True 值转换为 1)。

n = 3
pattern_value = True

res = df.with_columns(
    (pl.col('Pattern') == pattern_value).cast(pl.UInt8).rolling_sum(window_size=n)
)

对于 pattern_value = True 使用 (pl.col('Pattern') == pattern_value) 是多余的，您可以只使用 pl.col('模式').cast(pl.UInt8).

输出:

>>> res

shape: (7, 2)
┌────────────┬─────────┐
│ Date       ┆ Pattern │
│ ---        ┆ ---     │
│ str        ┆ u8      │
╞════════════╪═════════╡
│ 2023-01-01 ┆ null    │
│ 2023-01-02 ┆ null    │
│ 2023-01-03 ┆ 2       │
│ 2023-01-05 ┆ 2       │
│ 2023-01-10 ┆ 1       │
│ 2023-01-11 ┆ 1       │
│ 2023-01-12 ┆ 1       │
└────────────┴─────────┘

In pandas, this can be achieved by:

df["Pattern"].apply(lambda x: x == True).rolling(3, min_periods = 0).sum()

请注意，使用 min_periods = 0 不会产生上述输出。 Expr.rolling_sum如果需要，该方法还接受 min_periods 参数(默认为窗口大小)。例如

>>> df.with_columns(
    pl.col('Pattern').cast(pl.UInt8).rolling_sum(window_size=3, min_periods=0)
)

shape: (7, 2)
┌────────────┬─────────┐
│ Date       ┆ Pattern │
│ ---        ┆ ---     │
│ str        ┆ u8      │
╞════════════╪═════════╡
│ 2023-01-01 ┆ 1       │
│ 2023-01-02 ┆ 2       │
│ 2023-01-03 ┆ 2       │
│ 2023-01-05 ┆ 2       │
│ 2023-01-10 ┆ 1       │
│ 2023-01-11 ┆ 1       │
│ 2023-01-12 ┆ 1       │
└────────────┴─────────┘