如何为Polars DataFrame添加列存储每行值>0的列名列表？

针对你的需求，有两种简洁的实现方式：

方法一：列表推导式结合when/then #

遍历目标列，对每一列判断值是否大于0，满足条件时返回列名字符串，最后拼接成列表并去除空值：

import polars as pl

df = pl.DataFrame({"apple": [1, 0, 2, 0], "banana": [1, 0, 0, 1]})
cols = ["apple", "banana"]

result_df = df.with_columns(
    pl.concat_list(
        [pl.when(pl.col(col) > 0).then(col) for col in cols]
    ).list.drop_nulls().alias("result")
)

print(result_df)

方法二：利用struct与list.eval #

先将目标列转换为结构体，再通过列表筛选出值大于0的列名：

result_df = df.with_columns(
    pl.struct(cols)
    .list.eval(pl.element().filter(pl.element().value() > 0).keys())
    .alias("result")
)

print(result_df)

两种方法都能得到你期望的结果：

shape: (4, 3)
┌───────┬────────┬─────────────────────┐
│ apple ┆ banana ┆ result              │
│ ---   ┆ ---    ┆ ---                 │
│ i64   ┆ i64    ┆ list[str]           │
╞═══════╪════════╪═════════════════════╡
│ 1     ┆ 1      ┆ ["apple", "banana"] │
│ 0     ┆ 0      ┆ []                  │
│ 2     ┆ 0      ┆ ["apple"]           │
│ 0     ┆ 1      ┆ ["banana"]          │
└───────┴────────┴─────────────────────┘

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