如何用Polars原生API实现带当前行过滤的时间序列滚动窗口？

需求说明 #

针对金融交易Polars DataFrame，需要为每笔交易计算过去指定时长内的关联账户资金流入/流出总和：

• 获取当前行的source_acct和timestamp
• 回溯timestamp-d小时，筛选出source_acct或dest_acct与当前source_acct匹配的交易行
• 计算当前账户作为dest_acct时的金额总和（amount_in）
• 计算当前账户作为source_acct时的金额总和（amount_out，包含当前行）

现有实现问题 #

使用map_rows逐行计算时，在20M行数据集上速度极慢。尝试用Polars原生rollingAPI，但不知道如何结合账户维度的过滤逻辑。

原实现代码 #

from datetime import timedelta
import polars as pl

def windowing(df: pl.DataFrame, window_in_hours: int):
    d = timedelta(hours=window_in_hours)
    
    def calculate_amt(row):
        acc_no, window_end = row[0], row[1]
        window_start = window_end - d

        acct_window_mask = (
            (pl.col('timestamp') >= window_start) &
            (pl.col('timestamp') <= window_end) &
            (pl.col('dest_acct').eq(acc_no) | pl.col('source_acct').eq(acc_no))
        )
        window_txns = df.filter(acct_window_mask)

        amount_in = window_txns.filter(pl.col('dest_acct').eq(acc_no))['amount'].sum()
        amount_out = window_txns.filter(pl.col('source_acct').eq(acc_no))['amount'].sum()

        return (amount_in, amount_out)
    
    calculated_amounts = df.select(["source_acct", "timestamp", 'dest_acct', 'amount']).map_rows(calculate_amt)
    return df.with_columns(
        calculated_amounts['column_0'].alias('amount_in'), 
        calculated_amounts['column_1'].alias('amount_out'), 
    )

测试示例与预期输出 #

data = {
    "timestamp": [
        "2024-01-01 10:00:00",
        "2024-01-01 10:30:00",
        "2024-01-01 11:00:00",
        "2024-01-01 11:30:00",
        "2024-01-01 12:00:00"
    ],
    "source_acct": ["A", "B", "A", "C", "A"],
    "dest_acct": ["B", "A", "C", "A", "B"],
    "amount": [100, 150, 200, 300, 250]
}

df = pl.DataFrame(data).with_columns(pl.col("timestamp").str.to_datetime())
print(windowing(df, 1))

预期输出：

┌─────────────────────┬─────────────┬───────────┬────────┬───────────┬────────────┐
│ timestamp           ┆ source_acct ┆ dest_acct ┆ amount ┆ amount_in ┆ amount_out │
│ ---                 ┆ ---         ┆ ---       ┆ ---    ┆ ---       ┆ ---        │
│ datetime[μs]        ┆ str         ┆ str       ┆ i64    ┆ i64       ┆ i64        │
╞═════════════════════╪═════════════╪═══════════╪════════╪═══════════╪════════════╡
│ 2024-01-01 10:00:00 ┆ A           ┆ B         ┆ 100    ┆ 0         ┆ 100        │
│ 2024-01-01 10:30:00 ┆ B           ┆ A         ┆ 150    ┆ 100       ┆ 150        │
│ 2024-01-01 11:00:00 ┆ A           ┆ C         ┆ 200    ┆ 150       ┆ 300        │
│ 2024-01-01 11:30:00 ┆ C           ┆ A         ┆ 300    ┆ 200       ┆ 300        │
│ 2024-01-01 12:00:00 ┆ A           ┆ B         ┆ 250    ┆ 300       ┆ 450        │
└─────────────────────┴─────────────┴───────────┴────────┴───────────┴────────────┘

高效解决方案（Polars原生API实现） #

利用Polars的向量化分组、滚动窗口和连接操作，完全避免逐行循环，适合大规模数据集：

import polars as pl
from datetime import timedelta

def fast_windowing(df: pl.DataFrame, window_in_hours: int):
    # 确保数据按时间排序
    df_sorted = df.sort("timestamp")
    
    # 拆分交易记录：分别生成source账户的流出记录和dest账户的流入记录
    out_records = df_sorted.select(
        pl.col("timestamp"),
        pl.col("source_acct").alias("acct"),
        pl.col("amount").alias("out_amt"),
        pl.lit(0).alias("in_amt")
    )
    in_records = df_sorted.select(
        pl.col("timestamp"),
        pl.col("dest_acct").alias("acct"),
        pl.lit(0).alias("out_amt"),
        pl.col("amount").alias("in_amt")
    )
    
    # 合并记录并按账户、时间排序，为滚动窗口计算做准备
    combined = pl.concat([out_records, in_records]).sort(["acct", "timestamp"])
    
    # 按账户分组，计算指定时间窗口内的流入流出总和
    rolling_sums = combined.group_by("acct").rolling(
        index_column="timestamp",
        period=f"{window_in_hours}h",
        closed="both"
    ).agg(
        pl.col("in_amt").sum().alias("amount_in"),
        pl.col("out_amt").sum().alias("amount_out")
    )
    
    # 将滚动统计结果关联回原数据集
    result = df_sorted.join(
        rolling_sums,
        left_on=["source_acct", "timestamp"],
        right_on=["acct", "timestamp"],
        how="left"
    ).drop("acct")
    
    return result

# 测试
df = pl.DataFrame(data).with_columns(pl.col("timestamp").str.to_datetime())
print(fast_windowing(df, 1))

方案优势 #

1. 性能优化：全程使用Polars原生向量化操作，相比map_rows的逐行循环，性能提升几个数量级，可轻松处理20M行级别的数据集
2. 逻辑清晰：通过拆分source/dest记录，将所有与目标账户相关的交易统一到同一分组下，再用时间滚动窗口计算总和
3. 准确性保障：基于账户和时间的精准连接，确保每一行交易匹配到正确的窗口统计值

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