Finding (unique) set of values across columns - python-polars

Question

I have a Polars Dataframe that looks like the below
id
attribute
val1
val2
val3
1
True
A
A
A
2
True
A
A
B
2
False
A
B
C
I would like to create a new column that is the set of the values in val1, val2, and val3. For example,
id
attribute
val1
val2
val3
set
1
True
A
A
A
A
2
True
A
A
B
A, B
2
False
A
B
C
A, B, C
I can do something like this,
import polars as pl
df = pl.DataFrame({
'id': [1, 2, 3],
'attribute': [True, True, False],
'val1': ['A', 'A', 'A'],
'val2': ['A', 'A', 'B'],
'val3': ['A', 'B', 'C'],
})
df = df.with_columns([pl.struct('^val.*$').alias('set')])
df = df.with_columns(pl.col('set').apply(lambda x: set(x.values())))
However, with the apply, it is predictable slow. Is there a way to do this using native Polars functionality?

Answer 1

It's essentially a variation of https://stackoverflow.com/a/75387840/ as #ΩΠΟΚΕΚΡΥΜΜΕΝΟΣ pointed out.
df.with_columns(
pl.concat_list(pl.col("^val.*$"))
.arr.eval(
pl.element().unique(maintain_order=True).drop_nulls(),
parallel=True)
.alias("set")
)
shape: (3, 6)
┌─────┬───────────┬──────┬──────┬──────┬─────────────────┐
│ id | attribute | val1 | val2 | val3 | set │
│ --- | --- | --- | --- | --- | --- │
│ i64 | bool | str | str | str | list[str] │
╞═════╪═══════════╪══════╪══════╪══════╪═════════════════╡
│ 1 | true | A | A | A | ["A"] │
│ 2 | true | A | A | B | ["A", "B"] │
│ 3 | false | A | B | C | ["A", "B", "C"] │
└─────┴───────────┴──────┴──────┴──────┴─────────────────┘