如何在Polars中按组计算pct_change并匹配Pandas结果？

原始数据 #

首先定义初始的Polars DataFrame：

import polars as pl

df = pl.from_repr("""
┌───────┬─────┬──────┐
│ asset ┆ v   ┆ date │
│ ---   ┆ --- ┆ ---  │
│ str   ┆ i64 ┆ str  │
╞═══════╪═════╪══════╡
│ a     ┆ 1   ┆ 2017 │
│ b     ┆ 2   ┆ 2011 │
│ c     ┆ 3   ┆ 2012 │
│ a     ┆ 4   ┆ 2013 │
│ b     ┆ 5   ┆ 2014 │
│ c     ┆ 6   ┆ 2015 │
│ b     ┆ 7   ┆ 2016 │
│ c     ┆ 8   ┆ 2010 │
└───────┴─────┴──────┘
""")

需求说明 #

需要替代Pandas中groupby+apply的实现逻辑：按asset分组，每组内按date排序，然后计算v列的百分比变化（pct_change），最终得到如下结果：

┌───────┬─────┬──────┬────────────┐
│ asset ┆ v   ┆ date ┆ pct_change │
│ ---   ┆ --- ┆ ---  ┆ ---        │
│ str   ┆ i64 ┆ str  ┆ f64        │
╞═══════╪═════╪══════╪════════════╡
│ a     ┆ 4   ┆ 2013 ┆ null       │
│ a     ┆ 1   ┆ 2017 ┆ -0.75      │
│ b     ┆ 2   ┆ 2011 ┆ null       │
│ b     ┆ 5   ┆ 2014 ┆ 1.5        │
│ b     ┆ 7   ┆ 2016 ┆ 0.4        │
│ c     ┆ 8   ┆ 2010 ┆ null       │
│ c     ┆ 3   ┆ 2012 ┆ -0.625     │
│ c     ┆ 6   ┆ 2015 ┆ 1.0        │
└───────┴─────┴──────┴────────────┘

Polars实现方案 #

方案1：GroupBy + Apply（类Pandas逻辑） #

这种方式和Pandas的实现逻辑最接近，通过分组后应用自定义函数完成排序和计算：

def process_group(df: pl.DataFrame) -> pl.DataFrame:
    # 分组内按date排序，然后计算pct_change
    return df.sort("date").with_columns(
        pl.col("v").pct_change().alias("pct_change")
    )

# 按asset分组并应用函数
result = df.group_by("asset").apply(process_group)

方案2：窗口函数（更高效的矢量化实现） #

Polars支持在窗口函数中指定分组逻辑，结合全局排序后，能实现无Python层级循环的矢量化计算，性能更优：

# 先按asset和date全局排序，然后在asset分组内计算pct_change
result = df.sort(["asset", "date"]).with_columns(
    pl.col("v").pct_change().over("asset").alias("pct_change")
)

两种方案都能得到目标结果，其中方案2在处理大数据集时效率更高。

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