Polars pl.when().then().otherwise() в сочетании с первой строкой объекта group_by ⇐ Python

[Anonymous]

У меня есть pl.DataFrame с несколькими столбцами: level_0, символ, сигнал и сделка. В столбце «Сделка» просто указывается, покупать или продавать соответствующий символ («A» и «B»). Он вычисляется по("level_0", "symbol").
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")
)

shape: (16, 4)
┌─────────┬────────┬────────┬───────┐
│ level_0 ┆ symbol ┆ signal ┆ trade │
│ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ str ┆ i64 ┆ i64 │
╞═════════╪════════╪════════╪═══════╡
│ 0 ┆ A ┆ 1 ┆ null │
│ 0 ┆ A ┆ 0 ┆ -1 │
│ 0 ┆ A ┆ 1 ┆ 1 │
│ 0 ┆ A ┆ 1 ┆ null │
│ 0 ┆ B ┆ 0 ┆ null │
│ 0 ┆ B ┆ 1 ┆ 1 │
│ 0 ┆ B ┆ 1 ┆ null │
│ 0 ┆ B ┆ 0 ┆ -1 │
│ 1 ┆ A ┆ 0 ┆ null │
│ 1 ┆ A ┆ 0 ┆ null │
│ 1 ┆ A ┆ 0 ┆ null │
│ 1 ┆ A ┆ 1 ┆ 1 │
│ 1 ┆ B ┆ 1 ┆ null │
│ 1 ┆ B ┆ 1 ┆ null │
│ 1 ┆ B ┆ 0 ┆ -1 │
│ 1 ┆ B ┆ 0 ┆ null │
└─────────┴────────┴────────┴───────┘

Пока все хорошо. Единственное, первая строка каждой группы (["level_0", "symbol"]) неверно. Я хотел бы изменить нулевые значения в столбце trade в соответствии со следующим правилом:

• Если столбец сигнала содержит значение, отличное от 0 в соответствующей первой строке каждой группы, это значение должно быть скопировано в соответствующую первую строку каждой группы столбца сделки.
• Если столбец сигнала содержит равное значение до 0 в соответствующей первой строке каждой группы, значение соответствующей первой строки каждой группы столбца trade остается неизменным.

Иными словами, я хочу изменить нулевые значения в первой строке каждой группы в столбце trade всякий раз, когда в столбце сигнала есть значение, отличное от нуля.
Вот что я ищу:
shape: (16, 4)
┌─────────┬────────┬────────┬───────┐
│ level_0 ┆ symbol ┆ signal ┆ trade │
│ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ str ┆ i64 ┆ i64 │
╞═════════╪════════╪════════╪═══════╡
│ 0 ┆ A ┆ 1 ┆ 1 │

Подробнее здесь: https://stackoverflow.com/questions/791 ... roup-by-ob