如何在Polars中实现与Pandas一致的重复时间戳滚动计算结果？

问题场景 #

使用Pandas和Polars执行滚动求和操作时，遇到重复时间戳的场景下计算结果不一致：Pandas会为同一时间戳的每行单独计算滚动求和，而Polars默认会对同一时间戳的行做聚合处理，导致同一时间戳的结果完全相同。

Pandas实现及输出 #

import pandas as pd
import polars as pl

data = {
    "timestamp": [
        "2023-08-04 10:00:00",
        "2023-08-04 10:05:00",
        "2023-08-04 10:10:00",
        "2023-08-04 10:10:00",
        "2023-08-04 10:20:00",
        "2023-08-04 10:20:00",
    ],
    "value": [1, 2, 3, 4, 5, 6],
}

pddf = pd.DataFrame(data)
pddf["timestamp"] = pd.to_datetime(pddf["timestamp"])
pddf.set_index("timestamp", inplace=True)
pddf["rolling_sum"] = pddf["value"].rolling("10min").sum()
print(pddf)

输出：

value  rolling_sum
timestamp                              
2023-08-04 10:00:00      1          1.0
2023-08-04 10:05:00      2          3.0
2023-08-04 10:10:00      3          5.0
2023-08-04 10:10:00      4          9.0
2023-08-04 10:20:00      5          5.0
2023-08-04 10:20:00      6         11.0

Polars默认实现及输出 #

pldf = (
    pl.DataFrame(data)
    .with_columns(pl.col("timestamp").str.strptime(pl.Datetime))
    .sort("timestamp")
    .with_columns(
        pl.col("value")
        .rolling_sum_by(by="timestamp", window_size="10m")
        .alias("rolling_sum")
    )
)
print(pldf)

输出：

┌─────────────────────┬───────┬─────────────┐
│ timestamp           ┆ value ┆ rolling_sum │
│ ---                 ┆ ---   ┆ ---         │
│ datetime[μs]        ┆ i64   ┆ i64         │
╞═════════════════════╪═══════╪═════════════╡
│ 2023-08-04 10:00:00 ┆ 1     ┆ 1           │
│ 2023-08-04 10:05:00 ┆ 2     ┆ 3           │
│ 2023-08-04 10:10:00 ┆ 3     ┆ 9           │
│ 2023-08-04 10:10:00 ┆ 4     ┆ 9           │
│ 2023-08-04 10:20:00 ┆ 5     ┆ 11          │
│ 2023-08-04 10:20:00 ┆ 6     ┆ 11          │
└─────────────────────┴───────┴─────────────┘

解决方案：让Polars复刻Pandas行为 #

核心思路是让Polars能够逐行处理重复时间戳，而非先聚合相同时间戳的行。可以通过给重复时间戳添加微小的微秒偏移，确保每行拥有唯一的时间戳，再使用Polars的常规rolling方法计算：

pldf = (
    pl.DataFrame(data)
    .with_columns(pl.col("timestamp").str.strptime(pl.Datetime))
    .sort("timestamp")
    # 添加行号，为重复时间戳生成唯一的微秒偏移
    .with_row_index(name="idx")
    .with_columns(
        (pl.col("timestamp") + pl.duration(microseconds=pl.col("idx"))).alias("unique_ts")
    )
    # 基于唯一时间戳执行滚动求和，窗口范围10分钟
    .with_columns(
        pl.col("value")
        .rolling(
            index_column="unique_ts",
            window_size="10m",
            closed="both"
        )
        .sum()
        .alias("rolling_sum")
    )
    # 清理临时列
    .drop("idx", "unique_ts")
)
print(pldf)

输出结果将与Pandas完全一致：

┌─────────────────────┬───────┬─────────────┐
│ timestamp           ┆ value ┆ rolling_sum │
│ ---                 ┆ ---   ┆ ---         │
│ datetime[μs]        ┆ i64   ┆ i64         │
╞═════════════════════╪═══════╪═════════════╡
│ 2023-08-04 10:00:00 ┆ 1     ┆ 1           │
│ 2023-08-04 10:05:00 ┆ 2     ┆ 3           │
│ 2023-08-04 10:10:00 ┆ 3     ┆ 5           │
│ 2023-08-04 10:10:00 ┆ 4     ┆ 9           │
│ 2023-08-04 10:20:00 ┆ 5     ┆ 5           │
│ 2023-08-04 10:20:00 ┆ 6     ┆ 11          │
└─────────────────────┴───────┴─────────────┘

原理说明 #

• Pandas的rolling是逐行计算，窗口包含当前行及之前10分钟内的所有行（包括同一时间戳中排在当前行之前的行）。
• Polars的rolling_sum_by默认会先对相同by列（即timestamp）的行做聚合，再计算滚动窗口，因此同一时间戳的结果相同。
• 通过添加微秒偏移生成唯一时间戳后，使用Polars的rolling方法（指定index_column），就能实现和Pandas一致的逐行滚动计算逻辑。

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