0

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

var_1 var_2 var_3 var_4
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)
mark0512
  • 37
  • 3

1 Answers1

2

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