如何结合分组对象的首行使用Polars的pl.when().then().otherwise()方法？

嘿，看起来你正在用Polars处理分组后的信号转交易指令的逻辑，而且应该是卡在了分组首行的trade值处理上对吧？咱们用pl.when().then().otherwise()配合分组首行的信息来搞定这个问题～

首先先还原你的原始代码和当前的问题：你按level_0和symbol分组计算signal的差分得到trade列，但每个level_0+symbol分组的第一行因为没有前序值，diff()会返回null，这时候我们可以结合分组首行的signal值来补全这个空缺。

原始代码回顾 #

import polars as pl

pl.Config(tbl_rows=16)

df = pl.DataFrame(
    {
        "level_0": [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1],
        "symbol": [ "A", "A", "A", "A", "B", "B", "B", "B", "A", "A", "A", "A", "B", "B", "B", "B", ],
        "signal": [1, 0, 1, 1, 0, 1, 1, 0, 0, 0, 0, 1, 1, 1, 0, 0],
    }
).with_columns(
    pl.col("signal")
    .diff()
    .replace(old=0, new=None)
    .over("level_0", "symbol")
    .alias("trade")
)

此时每个分组的首行trade都是null，比如level_0=0, symbol=A的第一行，因为没有前序signal值，diff()无法计算。

解决方案：用pl.first()+条件判断补全首行 #

我们可以用pl.first("signal").over("level_0", "symbol")获取每个分组的首行signal值，再通过pl.when().then().otherwise()实现条件赋值：

import polars as pl

pl.Config(tbl_rows=16)

df = pl.DataFrame(
    {
        "level_0": [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1],
        "symbol": [ "A", "A", "A", "A", "B", "B", "B", "B", "A", "A", "A", "A", "B", "B", "B", "B", ],
        "signal": [1, 0, 1, 1, 0, 1, 1, 0, 0, 0, 0, 1, 1, 1, 0, 0],
    }
).with_columns(
    # 先计算原始的trade列
    pl.col("signal")
    .diff()
    .replace(old=0, new=None)
    .over("level_0", "symbol")
    .alias("trade")
).with_columns(
    # 针对分组首行的null值做条件赋值
    pl.when(pl.col("trade").is_null() & (pl.first("signal").over("level_0", "symbol") == 1))
    .then(pl.lit(1))  # 首行signal为1时，标记为开仓（比如买入，值可自定义）
    .when(pl.col("trade").is_null() & (pl.first("signal").over("level_0", "symbol") == 0))
    .then(pl.lit(-1)) # 首行signal为0时，标记为平仓（比如卖出，值可自定义）
    .otherwise(pl.col("trade"))  # 非首行的trade值保持不变
    .alias("trade")
)

代码解释 #

1. pl.first("signal").over(...)：这个表达式会为每一行返回其所属level_0+symbol分组的第一行signal值，相当于把分组首行的signal广播到整个分组的每一行。
2. 条件判断逻辑：
   • 如果当前行的trade是null（说明是分组首行）且分组首行signal=1，就把trade设为1（代表开仓，你可以根据业务逻辑改成其他值）
   • 如果当前行的trade是null且分组首行signal=0，就把trade设为-1（代表平仓，同样可自定义）
   • 其他情况（非首行的trade值）保持原样

灵活调整需求 #

如果你的业务逻辑里，首行signal=0时不需要标记（保持null），那可以简化条件：

.with_columns(
    pl.when(pl.col("trade").is_null() & (pl.first("signal").over("level_0", "symbol") == 1))
    .then(pl.lit(1))
    .otherwise(pl.col("trade"))
    .alias("trade")
)

这样就能完美结合分组首行的信息，用Polars的条件表达式完成列值的处理啦！

备注：内容来源于stack exchange，提问作者Andi