Polars高效实现：根据值拼接对应列名的最优方法

输入数据 #

期望输出 #

高效实现方案 #

避免使用map_elements的逐行循环，改用Polars原生的向量化操作，性能提升明显，以下两种方案均可：

方案1：使用pl.fold批量拼接

import polars as pl

df = pl.DataFrame({
    "a": [0, 1],
    "b": [1, 0],
    "c": [0, 0],
    "d": [1, 0],
    "e": [0, 0]
})

cols = df.columns
result_df = df.with_columns(
    pl.fold(
        acc=pl.lit(""),
        function=lambda acc, col: acc + pl.when(pl.col(col) == 1).then(col).otherwise(""),
        exprs=cols
    ).alias("RESULT")
)

方案2：使用列表拼接+空值过滤

import polars as pl

df = pl.DataFrame({
    "a": [0, 1],
    "b": [1, 0],
    "c": [0, 0],
    "d": [1, 0],
    "e": [0, 0]
})

cols = df.columns
result_df = df.with_columns(
    pl.concat_list([pl.when(pl.col(c) == 1).then(c) for c in cols])
    .list.drop_nulls()
    .str.concat("")
    .alias("RESULT")
)

方案优势 #

两种方案均基于Polars的向量化引擎批量处理数据，避免了map_elements带来的逐行Python循环开销，在数据量较大时性能差距会非常显著。

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