如何使用Polars按组内参考值归一化分组数据？

我需要用Polars完成分组条目值的归一化操作——将每组内的value值除以该组中reference_state为ref的对应value值，生成normalised列。

示例数据如下：

import polars as pl

df = pl.from_repr("""
┌──────────┬─────────────────┬───────┐
│ group_id ┆ reference_state ┆ value │
│ ---      ┆ ---             ┆ ---   │
│ i64      ┆ str             ┆ i64   │
╞══════════╪═════════════════╪═══════╡
│ 1        ┆ ref             ┆ 5     │
│ 1        ┆ a               ┆ 3     │
│ 1        ┆ b               ┆ 1     │
│ 2        ┆ ref             ┆ 4     │
│ 2        ┆ a               ┆ 8     │
│ 2        ┆ b               ┆ 2     │
└──────────┴─────────────────┴───────┘
""")

我已经用Pandas实现了该逻辑：

df = df.to_pandas()

for (i, x) in df.groupby("group_id"):
    ref_val = x.loc[x["reference_state"] == "ref"]["value"]
    df.loc[df["group_id"] == i, "normalised"] = x["value"] / ref_val.to_list()[0]
    
pl.from_pandas(df)

期望输出结果：

shape: (6, 4)
┌──────────┬─────────────────┬───────┬────────────┐
│ group_id ┆ reference_state ┆ value ┆ normalised │
│ ---      ┆ ---             ┆ ---   ┆ ---        │
│ i64      ┆ str             ┆ i64   ┆ f64        │
╞══════════╪═════════════════╪═══════╪════════════╡
│ 1        ┆ ref             ┆ 5     ┆ 1.0        │
│ 1        ┆ a               ┆ 3     ┆ 0.6        │
│ 1        ┆ b               ┆ 1     ┆ 0.2        │
│ 2        ┆ ref             ┆ 4     ┆ 1.0        │
│ 2        ┆ a               ┆ 8     ┆ 2.0        │
│ 2        ┆ b               ┆ 2     ┆ 0.5        │
└──────────┴─────────────────┴───────┴────────────┘

Polars实现方案 #

方法一：窗口函数（推荐）

利用Polars的窗口函数over，直接在分组内筛选参考值并完成计算，代码简洁高效：

result_df = df.with_columns(
    normalised=pl.col("value") / pl.col("value").filter(pl.col("reference_state") == "ref").over("group_id")
)

说明：pl.col("value").filter(...).over("group_id")会在每个group_id分组内，筛选出reference_state为ref的value值，该值会自动广播到当前组的所有行，之后用每行的value除以这个参考值，得到归一化结果。

方法二：分组提取参考值后合并

先提取每个组的参考值，再和原表关联计算，适合需要单独处理参考值的场景：

# 提取每个分组的参考值
ref_values = df.filter(pl.col("reference_state") == "ref").select(
    "group_id", pl.col("value").alias("ref_value")
)

# 关联原表并计算归一化列
result_df = df.join(ref_values, on="group_id").with_columns(
    normalised=pl.col("value") / pl.col("ref_value")
).drop("ref_value")

两种方法都能得到符合预期的输出结果。

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