Nested polars.col()

I would like to get value from column whose name is in another column.

In my table there are many columns, so doing this with pl.when().then() is not an option.

As input we have this dataframe:

and would like transform this with with_columns():

This (pseudo)code illustrates what I have in mind (but ofc. this doesn’t work)

df = pl.DataFrame({"foo" : [1, 2, 3, 4],
                   "foo_count" : [23, 45 ,234, 55],
                   "bar" : [4 ,6 ,9, 2],
                   "bar_count" : [43, 45 ,453, 67],
                   "baz": [5,1,15, 3],
                   "baz_count" : [64, 43 ,231, 94],
                   "largest" : ["baz", "bar", "baz", "foo"]})

df.with_columns(
    pl.col(f"{pl.col('largest')}_count").alias("largest_count")
)

>Solution :

Here is a simple approach. First, we create expressions for masked columns using a when-then expression. Then, we rely on pl.coalesce to combine the masked columns. Still there is a naive python loop over the unique elements in largest.

df.with_columns(
    pl.coalesce(
        pl.when(pl.col("largest") == col).then(pl.col(f"{col}_count"))
        for col in df.get_column("largest").unique()
    ).alias("largest_count")
)

shape: (4, 8)
┌─────┬───────────┬─────┬───────────┬─────┬───────────┬─────────┬───────────────┐
│ foo ┆ foo_count ┆ bar ┆ bar_count ┆ baz ┆ baz_count ┆ largest ┆ largest_count │
│ --- ┆ ---       ┆ --- ┆ ---       ┆ --- ┆ ---       ┆ ---     ┆ ---           │
│ i64 ┆ i64       ┆ i64 ┆ i64       ┆ i64 ┆ i64       ┆ str     ┆ i64           │
╞═════╪═══════════╪═════╪═══════════╪═════╪═══════════╪═════════╪═══════════════╡
│ 1   ┆ 23        ┆ 4   ┆ 43        ┆ 5   ┆ 64        ┆ baz     ┆ 64            │
│ 2   ┆ 45        ┆ 6   ┆ 45        ┆ 1   ┆ 43        ┆ bar     ┆ 45            │
│ 3   ┆ 234       ┆ 9   ┆ 453       ┆ 15  ┆ 231       ┆ baz     ┆ 231           │
│ 4   ┆ 55        ┆ 2   ┆ 67        ┆ 3   ┆ 94        ┆ foo     ┆ 55            │
└─────┴───────────┴─────┴───────────┴─────┴───────────┴─────────┴───────────────┘