pythonpython-polarspolars

Use an expression dictionary to calculate row wise based on a column in polars


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?


Solution

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