Как сделать условное масштабирование в полярах? ⇐ Python

[Anonymous]

У меня есть фрейм данных Polars, и я хочу масштабировать столбец в зависимости от того, является ли его значение положительным или отрицательным, масштабироваться на соответствующую сумму (положительную или отрицательную) и по другому столбцу.
Ниже приведен пример того, что я пробовал.

Код: Выделить всё

import polars as pl

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

foo_is_neg = pl.col("foo") < 0

df.with_columns(
pl.when(foo_is_neg)
.then(pl.col("foo") / pl.col("foo").filter(foo_is_neg).sum())
.otherwise(pl.col("foo") / pl.col("foo").filter(~foo_is_neg).sum())
.over("bar")
.alias("foo2")
)

df

После выполнения приведенного выше кода я получил следующий результат, который, похоже, и нужен.

Код: Выделить всё

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 │
└─────┴─────┴──────────┘

Есть ли канонический/лучший способ сделать это? Если да, то каким это должно быть?
Спасибо за помощь!

Подробнее здесь: https://stackoverflow.com/questions/752 ... -in-polars