[Python] How to do conditional scaling in polars?

I have a polars dataframe and I want to do scaling of a column depending on whether its value is positive or negative, scaled by respective sum (positive or negative) and over another column.

Below is an example for illustration for what I tried.

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


After running through the code above, I got the following output which seems to be what I want.

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


Is there a canonical/better way to do it? If so, what should it be like?

Thanks for your help!

Continue reading...