如何在Polars中结合非日期列用group_by_dynamic计算5年移动平均

我有一个Polars DataFrame，需要计算5年简单移动平均。但我不想仅对year列使用group_by_dynamic，还希望同时按地理单元State_District进行分组。

示例数据如下：

┌────────────┬─────────┬──────────┬────────────────────────────┬────────────┬────────────────────────────┬────────────────┬───────────┐
│ year       ┆ state   ┆ district ┆ candidate                  ┆ totalvotes ┆ Candidate_Name             ┆ State_District ┆ Last_Name │
│ ---        ┆ ---     ┆ ---      ┆ ---                        ┆ ---        ┆ ---                        ┆ ---            ┆ ---       │
│ str        ┆ str     ┆ i64      ┆ str                        ┆ i64        ┆ str                        ┆ str            ┆ str       │
╞════════════╪═════════╪══════════╪════════════════════════════╪════════════╪════════════════════════════╪════════════════╪═══════════╡
│ 1976-01-01 ┆ alabama ┆ 1        ┆ BILL DAVENPORT             ┆ 157170     ┆ bill davenport             ┆ alabama-1      ┆ davenport │
│ 1976-01-01 ┆ alabama ┆ 1        ┆ JACK EDWARDS               ┆ 157170     ┆ jack edwards               ┆ alabama-1      ┆ edwards   │
│ 1976-01-01 ┆ alabama ┆ 1        ┆ WRITEIN                    ┆ 157170     ┆ writein                    ┆ alabama-1      ┆ writein   │
│ 1976-01-01 ┆ alabama ┆ 2        ┆ J CAROLE KEAHEY            ┆ 156362     ┆ j carole keahey            ┆ alabama-2      ┆ keahey    │
│ 1976-01-01 ┆ alabama ┆ 2        ┆ WILLIAM L "BILL" DICKINSON ┆ 156362     ┆ william l "bill" dickinson ┆ alabama-2      ┆ dickinson │
└────────────┴─────────┴──────────┴────────────────────────────┴────────────┴────────────────────────────┴────────────────┴───────────┘

具体来说，我需要按year分组，并针对每个State_District计算5年简单移动平均。

我尝试过以下两种方法，但均未得到预期结果：

第一种方法：

mapped = election_lab.filter(
    pl.col("party") == "DEMOCRAT"
).group_by_dynamic(
    ["year", "State_District"], every = "5y"
).agg(
    pl.map_groups(exprs = ["candidatevotes", "totalvotes"], function = lambda x: x[0]/x[1]).alias("Dem_Vote_Share")
)

该方法报错，提示无法向group_by_dynamic传入列表。

第二种方法：

mapped = election_lab.filter(
    pl.col("party") == "DEMOCRAT"
).group_by(
    ["year", "State_District"]
).agg(
    pl.map_groups(exprs = ["candidatevotes", "totalvotes"], function = lambda x: x[0]/x[1]).alias("Dem_Vote_Share")
)

mapped_dynamic = mapped.group_by_dynamic(
    "year", every = "5y"    
).agg(
    pl.avg("Dem_Vote_Share").alias("SMA")
)

此方法返回全null的SMA列，也不符合需求。

请问最高效的实现方式是什么？

核心思路 #

要实现按State_District分组，同时对每个分组计算5年简单移动平均，需先通过普通group_by算出每个State_District每年的民主党得票率，再结合group_by_dynamic的group_by参数实现分组滚动计算。

正确代码实现 #

import polars as pl

# 第一步：计算每个State_District每年的民主党得票率
vote_share = election_lab.filter(pl.col("party") == "DEMOCRAT")
.group_by(["year", "State_District"])
.agg(
    # 用原生列运算替代map_groups，效率更高
    (pl.col("candidatevotes").sum() / pl.col("totalvotes").first()).alias("Dem_Vote_Share")
)
# 确保year转为日期类型（若原始为字符串）
.with_columns(pl.col("year").str.to_date())

# 第二步：按State_District分组，计算5年移动平均
sma_result = vote_share.group_by_dynamic(
    index_column="year",
    every="1y",  # 每年输出一个移动平均结果
    period="5y", # 窗口覆盖过去5年
    group_by="State_District"  # 关键：按地理单元分组计算
)
.agg(
    pl.col("Dem_Vote_Share").mean().alias("5y_SMA")
)

关键说明 #

1. 之前方法失败的原因：

   • 第一种方法错误地将分组列传入group_by_dynamic的第一个参数，该参数仅用于指定时间索引列，分组需通过group_by参数设置。
   • 第二种方法未指定group_by参数，导致group_by_dynamic对全局数据计算，而非按State_District单独计算，同时窗口设置不匹配数据时间间隔，最终返回null。

2. 效率优化点：

   • 用pl.col("candidatevotes").sum() / pl.col("totalvotes").first()替代map_groups，避免Python lambda的性能开销，Polars原生列运算效率更高。
   • 确保year为日期类型，group_by_dynamic需要时间类型的索引列才能正确处理时间窗口。

内容的提问来源于stack exchange，提问作者Damon C. Roberts