Create Polars list column of column names on each row that match a boolean condition
I have a simple pl.DataFrame with a number of columns that only contain boolean values.
import polars as pl
df = pl.DataFrame(
{"s1": [True, True, False], "s2": [False, True, True], "s3": [False, False, False]}
)
shape: (3, 3)
┌───────┬───────┬───────┐
│ s1 ┆ s2 ┆ s3 │
│ --- ┆ --- ┆ --- │
│ bool ┆ bool ┆ bool │
╞═══════╪═══════╪═══════╡
│ true ┆ false ┆ false │
│ true ┆ true ┆ false │
│ false ┆ true ┆ false │
└───────┴───────┴───────┘
I need to add another column that contains lists of varying length. A list in any individual row should contain the column name where the values of the columns S1, s2, and s3 have a True value.
Here's what I am actually looking for:
shape: (3, 4)
┌───────┬───────┬───────┬──────────────┐
│ s1 ┆ s2 ┆ s3 │ list │
│ --- ┆ --- ┆ --- │ --- │
│ bool ┆ bool ┆ bool │ list[str] │
╞═══════╪═══════╪═══════╡══════════════╡
│ true ┆ false ┆ false │ ["s1"] │
│ true ┆ true ┆ false │ ["s1", "s2"] │
│ false ┆ true ┆ false │ ["s2"] │
└───────┴───────┴───────┴──────────────┘