CODEWITHSUNDEEP
python-polars
mark0512
mark0512

Reputation: 77

Apply to a list of columns in Polars

In the following dataframe I would like to multiply var_3 and var_4 by negative 1. I can do so using the following method but I am wondering if it can be done by collecting them in a list (imagining that there may be many more than 4 columns in the dataframe)

df = pl.DataFrame({"var_1": ["a", "a", "b"],
                   "var_2": ["c", "d", "e"], 
                   "var_3": [1, 2, 3],
                   "var_4": [4, 5, 6]})

df.with_columns([pl.col("var_3") * -1, 
                  pl.col("var_4") * -1])

Which returns the desired dataframe

shape: (3, 4)
┌───────┬───────┬───────┬───────┐
│ var_1 ┆ var_2 ┆ var_3 ┆ var_4 │
│ ---   ┆ ---   ┆ ---   ┆ ---   │
│ str   ┆ str   ┆ i64   ┆ i64   │
╞═══════╪═══════╪═══════╪═══════╡
│ a     ┆ c     ┆ -1    ┆ -4    │
│ a     ┆ d     ┆ -2    ┆ -5    │
│ b     ┆ e     ┆ -3    ┆ -6    │
└───────┴───────┴───────┴───────┘

My try at it goes like this, though it is not applying the multiplication:

var_list = ["var_3", "var_4"]
pl_cols_var_list = [pl.col(k) for k in var_list]
df.with_columns(pl_cols_var_list * -1)

Upvotes: 2

Views: 1480

Answers (1)

user18559875
user18559875

Reputation:

You were close. You can provide your list of variable names (as strings) directly to the polars.col expression:

var_list = ["var_3", "var_4"]
df.with_columns(pl.col(var_list) * -1)

shape: (3, 4)
┌───────┬───────┬───────┬───────┐
│ var_1 ┆ var_2 ┆ var_3 ┆ var_4 │
│ ---   ┆ ---   ┆ ---   ┆ ---   │
│ str   ┆ str   ┆ i64   ┆ i64   │
╞═══════╪═══════╪═══════╪═══════╡
│ a     ┆ c     ┆ -1    ┆ -4    │
│ a     ┆ d     ┆ -2    ┆ -5    │
│ b     ┆ e     ┆ -3    ┆ -6    │
└───────┴───────┴───────┴───────┘

Another tip, if you have lots of columns and want to exclude only a few, you can use the polars.exclude expression:

var_list = ["var_1", "var_2"]
df.with_columns(pl.exclude(var_list) * -1)

shape: (3, 4)
┌───────┬───────┬───────┬───────┐
│ var_1 ┆ var_2 ┆ var_3 ┆ var_4 │
│ ---   ┆ ---   ┆ ---   ┆ ---   │
│ str   ┆ str   ┆ i64   ┆ i64   │
╞═══════╪═══════╪═══════╪═══════╡
│ a     ┆ c     ┆ -1    ┆ -4    │
│ a     ┆ d     ┆ -2    ┆ -5    │
│ b     ┆ e     ┆ -3    ┆ -6    │
└───────┴───────┴───────┴───────┘

Upvotes: 3

Related Questions