Alternatives to long chains of .when().then().when().then().otherwise()

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)
df.with_columns(
    pl.when(pl.col("Market") == "AT").then(pl.col("Number of Days") + params["AT"]["Value"]),
    pl.when(pl.col("Market") == "DE").then(pl.col("Number of Days") + params["DE"]["Value"]),
    pl.when(pl.col("Market") == "CA").then(pl.col("Number of Days") + params["CA"]["Value"]),
    pl.when(pl.col("Market") == "UK").then(pl.col("Number of Days") + params["UK"]["Value"]),
    pl.when(pl.col("Market") == "US").then(pl.col("Number of Days") + params["US"]["Value"])
)