pythonlist-comprehensionpython-polars

Use list comprehension for expression functions


Let's say I want to make a list of functions, ie aggs=['sum','std','mean','min','max']

then if I have an arbitrary df

df=pl.DataFrame({'a':[1,2,3], 'b':[2,3,4]})

I want to be able to do something like (this obviously doesn't work)

df.with_columns([pl.col('a').x() for x in aggs])

Is there a way to do that? aggs need not be a list of strings but just the easiest way to type out my intention for the purpose of this question. Additionally it'd need to have room for .name.suffix()

I know I could have a function that has all the aggs in the function and takes arbitrary dfs as a parameter which is like my backup plan so I'm hoping for something that resembles the above.


Solution

  • Would this work for you?

    df.with_columns([getattr(pl.col("a"), x)().name.suffix("_" + x) for x in aggs])
    
    shape: (3, 7)
    ┌─────┬─────┬───────┬───────┬────────┬───────┬───────┐
    │ a   ┆ b   ┆ a_sum ┆ a_std ┆ a_mean ┆ a_min ┆ a_max │
    │ --- ┆ --- ┆ ---   ┆ ---   ┆ ---    ┆ ---   ┆ ---   │
    │ i64 ┆ i64 ┆ i64   ┆ f64   ┆ f64    ┆ i64   ┆ i64   │
    ╞═════╪═════╪═══════╪═══════╪════════╪═══════╪═══════╡
    │ 1   ┆ 2   ┆ 6     ┆ 1.0   ┆ 2.0    ┆ 1     ┆ 3     │
    │ 2   ┆ 3   ┆ 6     ┆ 1.0   ┆ 2.0    ┆ 1     ┆ 3     │
    │ 3   ┆ 4   ┆ 6     ┆ 1.0   ┆ 2.0    ┆ 1     ┆ 3     │
    └─────┴─────┴───────┴───────┴────────┴───────┴───────┘