如何在Polars中按分组实现列值正负分别缩放？

你的现有实现逻辑是正确的，但可以从可读性和性能两方面进行优化，尤其是避免在when/then中重复执行filter操作（大数据量下会增加计算开销）。以下是两种更规范的实现方式：

方式一：提前计算分组内的正负总和（推荐） #

先通过窗口函数计算每个分组内的正值总和、负值总和，再基于这两个值进行缩放计算，代码逻辑更直观，且只需要一次分组统计：

import polars as pl

df = pl.DataFrame(
    {
        "foo": [1, 3, -2, -1, 1, 2],
        "bar": [1, 1, 1, 1, 2, 2],
    }
)

result = df.with_columns(
    # 计算分组内的正值总和（仅汇总foo>0的部分）
    pos_sum=pl.col("foo").filter(pl.col("foo") > 0).sum().over("bar"),
    # 计算分组内的负值总和（仅汇总foo<0的部分）
    neg_sum=pl.col("foo").filter(pl.col("foo") < 0).sum().over("bar")
).with_columns(
    foo2=pl.when(pl.col("foo") < 0)
          .then(pl.col("foo") / pl.col("neg_sum"))
          .otherwise(pl.col("foo") / pl.col("pos_sum"))
).drop(["pos_sum", "neg_sum"])

print(result)

方式二：使用group_by+join实现 #

如果偏好先分组计算统计量再合并的模式，可以用group_by生成分组统计结果，再通过join关联回原表：

import polars as pl

df = pl.DataFrame(
    {
        "foo": [1, 3, -2, -1, 1, 2],
        "bar": [1, 1, 1, 1, 2, 2],
    }
)

# 分组计算正负总和
group_stats = df.group_by("bar").agg(
    pos_sum=pl.col("foo").filter(pl.col("foo") > 0).sum(),
    neg_sum=pl.col("foo").filter(pl.col("foo") < 0).sum()
)

# 关联原表并计算缩放值
result = df.join(group_stats, on="bar").with_columns(
    foo2=pl.when(pl.col("foo") < 0)
          .then(pl.col("foo") / pl.col("neg_sum"))
          .otherwise(pl.col("foo") / pl.col("pos_sum"))
).drop(["pos_sum", "neg_sum"])

print(result)

优化说明 #

• 性能提升：两种优化方式都只执行一次分组内的正负总和计算，避免了原代码中when/then分支里重复调用filter+sum的冗余计算，大数据量下优势明显。
• 可读性增强：将统计逻辑和缩放逻辑分离，代码结构更清晰，后续维护更方便。
• 鲁棒性扩展：如果分组内只有正值或只有负值，filter会返回空序列，sum()会得到0，此时除法会报错。可以根据需求添加判断，比如：

  foo2=pl.when(pl.col("foo") < 0)
        .then(pl.col("foo") / pl.when(pl.col("neg_sum") != 0).then(pl.col("neg_sum")).otherwise(1))
        .otherwise(pl.col("foo") / pl.when(pl.col("pos_sum") != 0).then(pl.col("pos_sum")).otherwise(1))
  

输出结果与你的原代码完全一致：

shape: (6, 3)
┌─────┬─────┬──────────┐
│ foo ┆ bar ┆ foo2     │
│ --- ┆ --- ┆ ---      │
│ i64 ┆ i64 ┆ f64      │
╞═════╪═════╪══════════╡
│ 1   ┆ 1   ┆ 0.25     │
│ 3   ┆ 1   ┆ 0.75     │
│ -2  ┆ 1   ┆ 0.666667 │
│ -1  ┆ 1   ┆ 0.333333 │
│ 1   ┆ 2   ┆ 0.333333 │
│ 2   ┆ 2   ┆ 0.666667 │
└─────┴─────┴──────────┘

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