Polars 中获取每行最大值所在列名的完整教程

在 polars 中，可通过 `pl.coalesce` 结合 `pl.when/then` 动态匹配每行最大值对应的列名，实现类似 pandas `idxmax(axis=1)` 的功能，无需循环或 udf，完全向量化且高效。

要获取每行中最大值所在的列名（而非数值本身），Polars 并未提供直接等价于 Pandas idxmax(axis=1) 的内置表达式（如 pl.idxmax_horizontal 尚未存在），但可通过组合现有表达式优雅、高效地实现该逻辑。

核心思路是：

1. 先用 pl.max_horizontal("a", "b") 计算每行的最大值；
2. 对每个目标列，判断其值是否等于该行最大值；
3. 若相等，则返回对应列名（用 pl.lit(name)）；
4. 使用 pl.coalesce(...) 按顺序“回退选取”首个满足条件的列名（因每行仅有一个最大值，且 coalesce 保证返回第一个非-null结果）。

以下是完整可运行示例：

import polars as pl

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

# 获取每行最大值所在列名（字符串）
df = df.with_columns(
    max_col=pl.coalesce(
        pl.when(pl.col(name) == pl.max_horizontal(df.columns))
        .then(pl.lit(name))
        for name in df.columns
    )
)

print(df)

输出：

shape: (3, 3)
┌─────┬──────┬─────────┐
│ a   ┆ b    ┆ max_col │
│ --- ┆ ---  ┆ ---     │
│ i64 ┆ i64  ┆ str     │
╞═════╪══════╪═════════╡
│ 1   ┆ 4    ┆ b       │
│ 8   ┆ 5    ┆ a       │
│ 3   ┆ null ┆ a       │
└─────┴──────┴─────────┘

✅ 关键说明：

• pl.max_horizontal(df.columns) 自动作用于所有列（也可显式传入 ["a", "b"]）；
• pl.when(...).then(...) 返回一个惰性表达式，配合生成器推导式可批量构建多分支逻辑；
• pl.coalesce 按列顺序依次尝试，一旦某列满足条件即返回其列名，天然支持“首个最大列”语义（若多列并列最大，返回字典序最先者）；
• None 值会被 max_horizontal 自动忽略（符合 Polars 默认行为），因此 row=[3, null] 正确返回 "a"。

⚠️ 注意事项：

• 若需处理严格多列并列最大时返回全部列名（如 ["a", "b"] 列表），则需改用 list + apply（牺牲性能），不推荐；本方案适用于典型“取首个最优列”场景；
• 列名必须为合法标识符（无空格/特殊字符），否则 pl.lit(name) 仍有效，但后续操作需注意引用方式；
• 表达式完全惰性执行，零 Python 层循环，性能与原生 Polars 操作一致。

综上，该模式是 Polars 生态中实现 idxmax(axis=1) 语义的标准、高效、声明式解法，值得纳入日常数据处理工具箱。