I want to use an expression dictionary to perform calculations for a new column. I have this Polars dataframe:
df=pl.DataFrame(
"col1": ["a", "b", "a"],
"x": [1,2,3],
"y": [2,2,5]
)
And I have an expression dictionary:
expr_dict = {
"a": pl.col("x") * pl.col("y"),
"b": pl.col("x"),
}
I want to create a column where each value is calculated based on a key in in another column, but I do not know how. I want to hhave result like this:
>>> df.with_columns(r=pl.col("col1").apply(lambda x: expr_dict[X])
>>> shape: (3, 3)
┌──────┬─────┬─────┬─────┐
│ col1 ┆ x ┆ y ┆ r │
│ --- ┆ --- ┆ --- ┆ --- │
│ str ┆ i64 ┆ i64 ┆ i64 │
╞══════╪═════╪═════╪═════╡
│ a ┆ 1 ┆ 2 ┆ 2 │
│ b ┆ 2 ┆ 2 ┆ 4 │
│ a ┆ 3 ┆ 5 ┆ 15 │
└──────┴─────┴─────┴─────┘
Is this possible?
pl.when()
for conditional expression.pl.coalesce()
to combine conditional expressions together.df.with_columns(
r = pl.coalesce(
pl.when(pl.col.col1 == k).then(v)
for k, v in expr_dict.items()
)
)
shape: (3, 4)
┌──────┬─────┬─────┬─────┐
│ col1 ┆ x ┆ y ┆ r │
│ --- ┆ --- ┆ --- ┆ --- │
│ str ┆ i64 ┆ i64 ┆ i64 │
╞══════╪═════╪═════╪═════╡
│ a ┆ 1 ┆ 2 ┆ 2 │
│ b ┆ 2 ┆ 2 ┆ 2 │
│ a ┆ 3 ┆ 5 ┆ 15 │
└──────┴─────┴─────┴─────┘