在Polars滚动窗口聚合中获取当前行值及帧内索引的实现方案

嗨，我来分享一下在Polars里实现这个需求的具体方案，刚好我也碰到过类似的问题，咱们一步步来拆解：

先明确需求与初始数据 #

首先我们有这样的订单数据，需要按用户分组，基于日期做1周的滚动窗口聚合，同时要拿到当前行的order_id，以及当前行在对应滚动帧里的索引位置：

import polars as pl

df = pl.DataFrame(
    {
        "order_id": ["o01", "o02", "o03", "o04", "o10", "o11", "o12", "o13"],
        "customer_id": ["ca", "ca", "ca", "ca", "cb", "cb", "cb", "cb"],
        "date": [
            "2024-04-03",
            "2024-04-04",
            "2024-04-04",
            "2024-04-11",
            "2024-04-02",
            "2024-04-02",
            "2024-04-03",
            "2024-05-13",
        ],
    },
    schema_overrides={"date": pl.Date},
)

之前的尝试与问题 #

我一开始写的代码是这样的，想着用first()来取当前行的值，用int_range生成帧索引，但结果完全不对——frame_index全是0，current_order_id也都是窗口里的第一个值，根本不是当前行的：

(
    df.sort("customer_id", "date")
    .rolling(
        index_column="date",
        period="1w",
        offset="0d",
        closed="left",
        group_by="customer_id",
    )
    .agg(
        frame_index=pl.int_range(pl.len()).first(),
        current_order_id=pl.col("order_id").first(),
        orders=pl.col("order_id"),
    )
)

比如ca用户的第二个2024-04-04记录，我想要current_order_id是o03，frame_index是1，但结果却是o02和0，完全不符合预期。

想要的正确结果 #

最终我期望得到的输出应该是这样的：

customer_id date        frame_index current_order_id    orders
str         date        i64         str                 list[str]
"ca"        2024-04-03  0           "o01"               ["o01", "o02", "o03"]
"ca"        2024-04-04  0           "o02"               ["o02", "o03"]
"ca"        2024-04-04  1           "o03"               ["o02", "o03"]
"ca"        2024-04-11  0           "o04"               ["o04"]
"cb"        2024-04-02  0           "o10"               ["o10", "o11", "o12"]
"cb"        2024-04-02  1           "o11"               ["o10", "o11", "o12"]
"cb"        2024-04-03  0           "o12"               ["o12"]
"cb"        2024-05-13  0           "o13"               ["o13"]

两种可行的解决方案 #

后来我摸索出两种好用的方法，分享给你：

方法一：先生成滚动列表，再匹配当前行位置 #

这个方法最直观，先给每个行计算对应的滚动窗口订单列表，然后直接用原表的order_id去匹配它在列表里的索引，同时直接保留原表的order_id作为当前行的值：

(
    df.sort("customer_id", "date")
    # 先计算每个行对应的滚动窗口订单列表
    .with_columns(
        orders=pl.col("order_id").rolling(
            index_column="date",
            period="1w",
            offset="0d",
            closed="left",
            group_by="customer_id",
        ).agg(pl.col("order_id"))
    )
    # 计算当前order_id在滚动列表里的索引，同时保留当前行的order_id
    .with_columns(
        frame_index=pl.col("orders").list.index_of(pl.col("order_id")),
        current_order_id=pl.col("order_id")
    )
    # 调整列顺序到我们期望的样子
    .select("customer_id", "date", "frame_index", "current_order_id", "orders")
)

这个方法的核心是利用Polars滚动窗口聚合和原表行一一对应的特性，每个滚动列表刚好对应原表的那一行，所以用list.index_of就能精准找到当前行在窗口里的位置。

方法二：在聚合步骤中直接关联当前行值 #

如果你的需求是要在agg步骤里用当前行的值做计算（比如和窗口均值做加减），可以先把当前行的order_id标记成一个独立列，然后在聚合时保留这个值，再用它去匹配帧内索引：

(
    df.sort("customer_id", "date")
    # 先标记当前行的order_id
    .with_columns(current_order_id=pl.col("order_id"))
    .rolling(
        index_column="date",
        period="1w",
        offset="0d",
        closed="left",
        group_by="customer_id",
    )
    .agg(
        orders=pl.col("order_id"),
        # 因为每个窗口对应原表一行，所以first()取的就是当前行的current_order_id
        current_order_id=pl.col("current_order_id").first(),
        # 用current_order_id匹配它在orders列表里的索引
        frame_index=pl.col("order_id").list.index_of(pl.col("current_order_id").first())
    )
    # 调整列顺序
    .select("customer_id", "date", "frame_index", "current_order_id", "orders")
)

这个方法同样能得到预期结果，适合需要在聚合过程中直接使用当前行值做计算的场景。

总结一下 #

其实核心就是抓住Polars滚动窗口的一个特性：每个滚动窗口聚合结果和原表的行是一一对应的，基于这个特性，我们可以轻松关联当前行的值，并用列表索引方法找到它在窗口内的位置。

备注：内容来源于stack exchange，提问作者dpprdan