
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) 的推荐实践。










