Sum columns based on column names in a list for polars

Viewed 124

So in python Polars

I can add one or more columns to make a new column by using an expression something like

frame.with_column((pl.col('colname1') + pl.col('colname2').alias('new_colname')))

However, if I have all the column names in a list is there a way to sum all the columns in that list and create a new column based on the result ?

Thanks

1 Answers

sum expr supports horizontal summing. From the docs,

List[Expr] -> aggregate the sum value horizontally.

Sample code for ref,

import polars as pl

df = pl.DataFrame({"a": [1, 2, 3], "b": [1, 2, None]})
print(df)

This results in something like,

shape: (3, 2)
┌─────┬──────┐
│ a   ┆ b    │
│ --- ┆ ---  │
│ i64 ┆ i64  │
╞═════╪══════╡
│ 1   ┆ 1    │
├╌╌╌╌╌┼╌╌╌╌╌╌┤
│ 2   ┆ 2    │
├╌╌╌╌╌┼╌╌╌╌╌╌┤
│ 3   ┆ null │
└─────┴──────┘

On this you can do something like,

cols = ["a", "b"]
df2 = df.select(pl.sum([pl.col(i) for i in cols]).alias('new_colname'))

print(df2)

Which will result in,

shape: (3, 1)
┌──────┐
│ sum  │
│ ---  │
│ i64  │
╞══════╡
│ 2    │
├╌╌╌╌╌╌┤
│ 4    │
├╌╌╌╌╌╌┤
│ null │
└──────┘
Related