如何在 Polars 中获取每行最大值所在列的列名

polars 本身不直接提供类似 pandas `idxmax(axis=1)` 的横向列名定位功能，但可通过 `pl.when().then()` 配合 `pl.coalesce()` 和 `pl.max_horizontal()` 组合实现——即逐列判断是否等于该行最大值，并返回对应列名。

在 Polars 中，pl.max_horizontal("a", "b") 能高效计算每行多列中的最大数值，但若目标是获取最大值所在的列名（而非值本身），就需要借助条件表达式构建逻辑映射。核心思路是：对每一列，判断其值是否等于该行在指定列集合中的最大值；若成立，则返回该列名；最后用 pl.coalesce() 将多个条件结果“合并”为单一字符串列，优先取首个非空匹配。

以下是一个完整、可复用的实现示例：

import polars as pl

df = pl.DataFrame(
    {
        "a": [1, 8, 3],
        "b": [4, 5, None],
    }
)

# 指定参与比较的列（支持任意数量）
target_cols = ["a", "b"]

df = df.with_columns(
    max_col=pl.coalesce(
        [
            pl.when(pl.col(name) == pl.max_horizontal(target_cols))
            .then(pl.lit(name))
            for name in target_cols
        ]
    )
)

print(df)

输出结果为：

shape: (3, 3)
┌─────┬──────┬─────────┐
│ a   ┆ b    ┆ max_col │
│ --- ┆ ---  ┆ ---     │
│ i64 ┆ i64  ┆ str     │

╞═════╪══════╪═════════╡
│ 1   ┆ 4    ┆ b       │
│ 8   ┆ 5    ┆ a       │
│ 3   ┆ null ┆ a       │
└─────┴──────┴─────────┘

✅ 关键说明：

• pl.max_horizontal(target_cols) 在每行内动态计算最大值（自动忽略 null，与 Pandas skipna=True 行为一致）；
• pl.when(...).then(pl.lit(name)) 生成一个惰性表达式，仅当条件为真时输出列名字符串；
• pl.coalesce() 按顺序尝试各表达式，返回第一个非-null 结果，确保每行仅有一个列名被选中（即使多列并列最大，也按列表顺序取首个匹配项）；
• 此方法天然支持 null 值处理，且完全向量化，无需 .apply() 或 Python 循环，性能优异。

⚠️ 注意事项：

• 若需处理严格并列最大值时返回所有列名（如 "a,b"），则需改用 pl.concat_list() + pl.list.eval() 方式，复杂度上升；
• 列名必须为合法标识符（避免空格或特殊字符），否则 pl.lit(name) 仍有效，但后续操作可能受限；
• target_cols 应预先校验存在性，防止 KeyError，生产环境建议添加 assert all(col in df.columns for col in target_cols)。

该方案是 Polars 社区广泛采用的惯用模式，兼顾简洁性、可读性与执行效率，是替代 Pandas idxmax(axis=1) 的推荐实践。