Can you create multiple columns based on the same set of conditions in Polars?

Is it possible to do something like this in Polars? Like do you need a separate when.then.otherwise for each of the 4 new varialbles, or can you use struct to create multiple new variables from one when.then.otherwise?

Regular Python example:

if x=1 and y=3 and w=300*z and z<100:
     tot = 300
     work = 400
     sie = 500
     walk = 'into'
else:
     tot = 350
     work = 400*tot
     sie = tot/1000
     walk = 'outof'

I tried to do a similar thing in Polars with struct (to create new variables a and b based on Movie variable:

import polars as pl

ratings = pl.DataFrame(
    {
        "Movie": ["Cars", "IT", "ET", "Cars", "Up", "IT", "Cars", "ET", "Up", "Cars"],
        "Theatre": ["NE", "ME", "IL", "ND", "NE", "SD", "NE", "IL", "IL", "NE"],
        "Avg_Rating": [4.5, 4.4, 4.6, 4.3, 4.8, 4.7, 4.5, 4.9, 4.7, 4.6],
        "Count": [30, 27, 26, 29, 31, 28, 28, 26, 33, 28],
    }
)

x = ratings.with_columns(
    pl.when(pl.col('Movie')=='Up').then(pl.struct(pl.lit(0),pl.lit(2))).otherwise(pl.struct(pl.lit(1),pl.lit(3))).struct.field(['a','b'])
)
print(x)

Thanks!

If you remove the .struct.field() call you will see the issue.

# DuplicateError: multiple fields with name 'literal' found

You need to give names to the struct fields you are creating.

df.with_columns(
    pl.when(pl.col('Movie') == 'Up')
      .then(pl.struct(pl.lit(0).alias('a'), pl.lit(2).alias('b')))
      .otherwise(pl.struct(pl.lit(1).alias('a'), pl.lit(3).alias('b')))
      .struct.field('a', 'b')
)

There are also some ways to neaten it up if you prefer.

• pl.when() uses kwargs as shorthand for equality conditions.
• pl.struct() uses kwargs as shorthand for .alias() to name fields.
• pl.lit() is not required for integers
• .struct.unnest() will unnest all fields.

df.with_columns(
    pl.when(Movie='Up')
      .then(pl.struct(a=0, b=2))
      .otherwise(pl.struct(a=1, b=3))
      .struct.unnest()
)

shape: (10, 6)
┌───────┬─────────┬────────────┬───────┬─────┬─────┐
│ Movie ┆ Theatre ┆ Avg_Rating ┆ Count ┆ a   ┆ b   │
│ ---   ┆ ---     ┆ ---        ┆ ---   ┆ --- ┆ --- │
│ str   ┆ str     ┆ f64        ┆ i64   ┆ i32 ┆ i32 │
╞═══════╪═════════╪════════════╪═══════╪═════╪═════╡
│ Cars  ┆ NE      ┆ 4.5        ┆ 30    ┆ 1   ┆ 3   │
│ IT    ┆ ME      ┆ 4.4        ┆ 27    ┆ 1   ┆ 3   │
│ ET    ┆ IL      ┆ 4.6        ┆ 26    ┆ 1   ┆ 3   │
│ Cars  ┆ ND      ┆ 4.3        ┆ 29    ┆ 1   ┆ 3   │
│ Up    ┆ NE      ┆ 4.8        ┆ 31    ┆ 0   ┆ 2   │
│ IT    ┆ SD      ┆ 4.7        ┆ 28    ┆ 1   ┆ 3   │
│ Cars  ┆ NE      ┆ 4.5        ┆ 28    ┆ 1   ┆ 3   │
│ ET    ┆ IL      ┆ 4.9        ┆ 26    ┆ 1   ┆ 3   │
│ Up    ┆ IL      ┆ 4.7        ┆ 33    ┆ 0   ┆ 2   │
│ Cars  ┆ NE      ┆ 4.6        ┆ 28    ┆ 1   ┆ 3   │
└───────┴─────────┴────────────┴───────┴─────┴─────┘