使用polars的滚动函数得到滚动中所有值的列表windows

Question

我想使用滚动函数来获取滚动中所有值的列表 windows。

我用下面的代码片段试了一下：

import polars as pl
df = pl.DataFrame(
    {
        "A": [1.0, 2.0, 9.0, 2.0, 13.0],
    }
)

df.select(
    [
        pl.col("A").rolling_apply(3, lambda s: s),
    ]
)

这个输出

┌──────┐
│ A    │
│ ---  │
│ f64  │
╞══════╡
│ null │
├╌╌╌╌╌╌┤
│ null │
├╌╌╌╌╌╌┤
│ 1    │
├╌╌╌╌╌╌┤
│ 2    │
├╌╌╌╌╌╌┤
│ 9    │
└──────┘

但我需要的是：

┌─────────────────┐
│ A               │
│ ---             │
│ list [f64]      │
╞═════════════════╡
│ [null, null, 1] │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ [null, 1, 2]    │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ [1, 2, 9]       │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ [2, 9, 2]       │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ [9, 2, 13]      │
└─────────────────┘

有没有人知道如何在极地中以简单的方式做到这一点？

Answer 1

您看到的行为是 lambda s: s 导致 rolling_apply 选择最后一个元素。目前似乎没有办法让它 return 多个元素。例如，当我尝试使用 lambda s: list(s) 将其强制添加到列表时，它会引发一个错误，即列表无法 returned，表明此功能不存在。

下面是涉及 Python 循环（太慢）的解决方法，但可能对您有所帮助：

定义自定义 slice 函数：

def _slice(s: pl.Series, offset: int, l: int) -> pl.Series:
    # like s.slice(offset, l), but prepadding with null when offset <0
    if offset < 0:
        prepad = pl.Series([None] * abs(offset))
        return pl.concat((prepad, s.slice(0, l+offset)))
    else:
        return s.slice(offset, l)

并用它来循环 df["A"]:

l = 3
pl.Series([_slice(df["A"], n-l+1, l) for n in range(len(df))])

这导致：

shape: (5,)
Series: '' [list]
[
    [null, null, 1]
    [null, 1, 2]
    [1, 2, 9]
    [2, 9, 2]
    [9, 2, 13]
]

这绝不是一个高效的解决方案，所以这里的问题是您将如何处理数据框中的列表？以这种格式存储数据对于进一步处理通常是一个挑战。

Answer 2

您可以创建滞后列并将它们收集到列表中。

(df
    .with_columns([pl.col("A").shift(i).alias(f"A_lag_{i}") for i in range(3)])
    .select(
        [pl.concat_list([f"A_lag_{i}" for i in range(3)][::-1]).alias("A_rolling")]
))

输出：

shape: (5, 1)
┌─────────────────┐
│ A_rolling       │
│ ---             │
│ list [f64]      │
╞═════════════════╡
│ [null, null, 1] │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ [null, 1, 2]    │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ [1, 2, 9]       │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ [2, 9, 2]       │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ [9, 2, 13]      │
└─────────────────┘

让我们分解一下：

reshape(-1, 1) 将 "A" 转换为列表。

pl.col(..).shift(i) for i in range(3)] 创建新的滞后列。

这导致了这个中间 DataFrame：

shape: (5, 4)
┌─────┬────────────┬────────────┬────────────┐
│ A   ┆ A_lag_0    ┆ A_lag_1    ┆ A_lag_2    │
│ --- ┆ ---        ┆ ---        ┆ ---        │
│ f64 ┆ list [f64] ┆ list [f64] ┆ list [f64] │
╞═════╪════════════╪════════════╪════════════╡
│ 1   ┆ [1]        ┆ [null]     ┆ [null]     │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 2   ┆ [2]        ┆ [1]        ┆ [null]     │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 9   ┆ [9]        ┆ [2]        ┆ [1]        │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 2   ┆ [2]        ┆ [9]        ┆ [2]        │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 13  ┆ [13]       ┆ [2]        ┆ [9]        │
└─────┴────────────┴────────────┴────────────┘

最后我们以相反的顺序连接它们并将输出命名为“A_rolling”:

pl.concat_list([f"A_lag_{i}" for i in range(3)][::-1]).alias("A_rolling")

使用polars的滚动函数得到滚动中所有值的列表windows

Use the rolling function of polars to get a list of all values in the rolling windows

python

python-polars