We want to avoid hard-coding the countries in the .with_columns()
when writing a long when().then().otherwise()
chain.
Let’s use the following dataframe and dictionary as an example:
df = pl.DataFrame(
{
"Market":["AT", "AT", "DE", "DE", "CA", "DE", "UK", "US"],
"Number of Days":[1, 2, 3, 4, 3, 4, 2, 1],
}
)
params = {
"AT":{"Value": 1},
"DE":{"Value": 2},
"CA":{"Value": 3},
"UK":{"Value": 1},
"US":{"Value": 2}
}
We can create a list of expressions based on the values provided in the dictionary, without hard-coding the countries in the .with_columns()
:
exprs = []
for market, data in params.items():
condition = (pl.col("Market") == market)
result = (pl.col("Number of Days") + data["Value"])
expr = pl.when(condition).then(result)
exprs.append(expr)
df.with_columns(exprs)