如何在Polars DataFrame中按组根据signal列首行值替换trade列单元格

问题背景 #

现有如下pl.DataFrame：

import polars as pl

df = pl.DataFrame(
    {
        "symbol": ["s1", "s1", "s2", "s2"],
        "signal": [0, 1, 2, 0],
        "trade": [None, 1, None, -1],
    }
)

其结构如下：

shape: (4, 3)
┌────────┬────────┬───────┐
│ symbol ┆ signal ┆ trade │
│ ---    ┆ ---    ┆ ---   │
│ str    ┆ i64    ┆ i64   │
╞════════╪════════╪═══════╡
│ s1     ┆ 0      ┆ null  │
│ s1     ┆ 1      ┆ 1     │
│ s2     ┆ 2      ┆ null  │
│ s2     ┆ 0      ┆ -1    │
└────────┴────────┴───────┘

需求：按symbol列分组，检查每组signal列的第一行值是否不等于0。若满足条件，则将该组第一行的trade列值替换为对应行的signal值。

期望得到的结果：

shape: (4, 3)
┌────────┬────────┬───────┐
│ symbol ┆ signal ┆ trade │
│ ---    ┆ ---    ┆ ---   │
│ str    ┆ i64    ┆ i64   │
╞════════╪════════╪═══════╡
│ s1     ┆ 0      ┆ null  │
│ s1     ┆ 1      ┆ 1     │
│ s2     ┆ 2      ┆ 2     │ <- 从signal列复制值
│ s2     ┆ 0      ┆ -1    │
└────────┴────────┴───────┘

解决方案 #

方法一：窗口函数实现（推荐，性能更优） #

直接利用Polars的窗口函数over获取分组内的首个signal值，再通过条件表达式替换trade列：

import polars as pl

df = pl.DataFrame(
    {
        "symbol": ["s1", "s1", "s2", "s2"],
        "signal": [0, 1, 2, 0],
        "trade": [None, 1, None, -1],
    }
)

result = df.with_columns(
    pl.when(pl.col("signal").first().over("symbol") != 0)
      .then(pl.col("signal"))
      .otherwise(pl.col("trade"))
      .alias("trade")
)

print(result)

方法二：分组映射实现 #

通过group_by结合map_groups遍历每个分组，在组内处理逻辑：

import polars as pl

df = pl.DataFrame(
    {
        "symbol": ["s1", "s1", "s2", "s2"],
        "signal": [0, 1, 2, 0],
        "trade": [None, 1, None, -1],
    }
)

result = df.group_by("symbol", maintain_order=True).map_groups(
    lambda group: group.with_columns(
        pl.when(group["signal"][0] != 0)
          .then(group["signal"])
          .otherwise(group["trade"])
          .alias("trade")
    )
)

print(result)

思路说明 #

1. 窗口函数方式：

   • 使用pl.col("signal").first().over("symbol")获取每个symbol分组内的第一个signal值。
   • 通过pl.when().then().otherwise()构建条件：如果分组首个signal不等于0，就用当前行的signal值替换trade；否则保留原trade值。
   • 这种方式无需遍历分组，性能更适合大数据场景。

2. 分组映射方式：

   • group_by("symbol", maintain_order=True)按symbol分组并保留原顺序。
   • map_groups传入匿名函数处理每个分组，判断分组首个signal值后替换trade列。

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