如何在Polars中按布尔条件生成含匹配列名的列表列

我有一个仅包含布尔值列的Polars DataFrame：

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

需要添加一列，每行的列表包含对应行中值为True的列名，期望结果如下：

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

实现方法 #

方法一：结构体+列表求值（简洁高效）

把所有列转为结构体，遍历元素筛选出值为True的列名，最后过滤空值：

result = df.with_columns(
    pl.struct(pl.all())
    .list.eval(pl.when(pl.element()).then(pl.element().name))
    .filter(pl.element().is_not_null())
    .alias("list")
)

方法二：数据重塑+分组聚合（适配大表场景）

先将宽表转长表，筛选True的行后按原行索引分组聚合列名：

result = (
    df.with_row_index("idx")
    .melt(id_vars="idx", variable_name="col", value_name="val")
    .filter(pl.col("val"))
    .group_by("idx")
    .agg(pl.col("col").alias("list"))
    .join(df.with_row_index("idx"), on="idx")
    .drop("idx")
    .select(df.columns + ["list"])
)

运行任意一种方法，都能得到预期结果。

内容的提问来源于stack exchange，提问作者Andi