如何高效对Polars DataFrame嵌套列表执行滑动窗口求和取最大值？

问题背景 #

给定如下Polars DataFrame：

import polars as pl
from polars import Boolean, List, col

src = pl.DataFrame(
    {
        "c1": ["a", "b", "c", "d"],
        "c2": [
            [0, 0],
            [0, 1, 0, 0],
            [1, 0, 1, 1, 1],
            [1, 1, 0],
        ],
    },
    schema_overrides={"c2": List(Boolean)},
)

需求为对c2列的每个内部列表，计算窗口长度为3的滑动窗口求和的最大值，逻辑示例：

[1, 0, 0, 1, 1, 1]
⇒
[(1, 0, 0), (0, 0, 1), (0, 1, 1),(1, 1, 1)]
⇒
[1, 1, 2, 3]
⇒
3

尝试过展开嵌套列表再执行滚动计算的方法，但该方法会导致滚动窗口跨越原内部列表边界，不符合需求，错误代码及输出如下：

(
    src.with_row_index("outer_index")
    .explode("c2")
    .with_row_index("inner_index")
    .rolling("inner_index", period="3i", offset="0i", closed="left")
    .agg("outer_index", "c1", col("c2").sum())
    .drop("inner_index")
)

错误输出（字典形式）：

[{'outer_index': [0, 0, 1], 'c1': ['a', 'a', 'b'], 'c2': 0},
 {'outer_index': [0, 1, 1], 'c1': ['a', 'b', 'b'], 'c2': 1},
 {'outer_index': [1, 1, 1], 'c1': ['b', 'b', 'b'], 'c2': 1},
 {'outer_index': [1, 1, 1], 'c1': ['b', 'b', 'b'], 'c2': 1},
 {'outer_index': [1, 1, 2], 'c1': ['b', 'b', 'c'], 'c2': 1},
 {'outer_index': [1, 2, 2], 'c1': ['b', 'c', 'c'], 'c2': 1},
 {'outer_index': [2, 2, 2], 'c1': ['c', 'c', 'c'], 'c2': 2},
 {'outer_index': [2, 2, 2], 'c1': ['c', 'c', 'c'], 'c2': 2},
 {'outer_index': [2, 2, 2], 'c1': ['c', 'c', 'c'], 'c2': 3},
 {'outer_index': [2, 2, 3], 'c1': ['c', 'c', 'd'], 'c2': 3},
 {'outer_index': [2, 3, 3], 'c1': ['c', 'd', 'd'], 'c2': 3},
 {'outer_index': [3, 3, 3], 'c1': ['d', 'd', 'd'], 'c2': 2},
 {'outer_index': [3, 3], 'c1': ['d', 'd'], 'c2': 1},
 {'outer_index': [3], 'c1': ['d'], 'c2': 0}]

高效解决方案 #

方法1：map_batches结合numpy向量化计算 #

利用numpy的sliding_window_view实现高效滑动窗口求和，适合大体积数据场景：

import numpy as np

def rolling_max_sum(arr: np.ndarray, window_size: int) -> int:
    if len(arr) < window_size:
        return arr.sum()  # 可根据需求调整短列表的返回逻辑，比如返回0
    windows = np.lib.stride_tricks.sliding_window_view(arr, window_size)
    return windows.sum(axis=1).max()

result = src.with_columns(
    rolling_max=pl.col("c2").map_batches(
        lambda s: s.struct.arr.to_numpy().apply(rolling_max_sum, window_size=3)
    )
)

print(result)

输出结果：

shape: (4, 3)
┌─────┬───────────────────┬────────────┐
│ c1  ┆ c2                ┆ rolling_max│
│ --- ┆ ---               ┆ ---        │
│ str ┆ list[bool]        ┆ i64        │
╞═════╪═══════════════════╪════════════╡
│ a   ┆ [false, false]    ┆ 0          │
│ b   ┆ [false, true, …]  ┆ 1          │
│ c   ┆ [true, false, …]  ┆ 3          │
│ d   ┆ [true, true, false]┆ 2          │
└─────┴───────────────────┴────────────┘

方法2：纯Polars列表操作（无外部依赖） #

通过列表切片生成滑动窗口，直接计算求和最大值，逻辑直观：

window_size = 3

result = src.with_columns(
    rolling_max=pl.col("c2").map_elements(
        lambda lst: max(
            sum(lst[i:i+window_size]) for i in range(len(lst) - window_size + 1)
        ) if len(lst) >= window_size else sum(lst),
        return_dtype=pl.Int64
    )
)

print(result)

该方法无需额外依赖，但对于极长列表，性能略低于numpy向量化方案。

方法3：修正分组滚动逻辑 #

如果偏好展开-滚动的思路，可通过按原行分组避免窗口跨边界：

window_size = 3

result = (
    src.with_row_index("idx")
    .explode("c2")
    .with_columns(inner_idx=pl.int_range(0, pl.count()).over("idx"))
    .rolling(
        index_column="inner_idx",
        window_size=window_size,
        by="idx",  # 限制窗口仅在原行内部滚动
        closed="left"
    )
    .agg(pl.col("c2").sum())
    .group_by("idx", "c1", maintain_order=True)
    .agg(pl.col("c2").max().alias("rolling_max"))
    .drop("idx")
)

print(result)

该方法通过by="idx"参数确保滚动窗口不会跨越原列表边界，同时保留Polars的向量化性能。

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