1

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

我用以下代码片段尝试了它：

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]      │
└─────────────────┘

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

4

2 回答 2

1

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

下面是一个涉及 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]
]

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

于 2021-12-17T11:06:07.307 回答
1

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

(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)]创建新的滞后列。

这导致了这个中间数据帧：

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

于 2021-12-18T19:07:27.133 回答