
在 polars 中,可通过 `pl.coalesce` 结合 `pl.when/then` 动态匹配每行最大值对应的列名,实现类似 pandas `idxmax(axis=1)` 的功能,无需循环或 udf,完全向量化且高效。
要获取每行中最大值所在的列名(而非数值本身),Polars 并未提供直接等价于 Pandas idxmax(axis=1) 的内置表达式(如 pl.idxmax_horizontal 尚未存在),但可通过组合现有表达式优雅、高效地实现该逻辑。
核心思路是:
- 先用 pl.max_horizontal("a", "b") 计算每行的最大值;
- 对每个目标列,判断其值是否等于该行最大值;
- 若相等,则返回对应列名(用 pl.lit(name));
- 使用 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) 语义的标准、高效、声明式解法,值得纳入日常数据处理工具箱。










