如何在Polars中获取每行最大值对应的列名？

在Polars中，我们可以通过pl.max_horizontal快速获取指定列集合中每行的最大值，对应Pandas里的df[["a", "b"]].max(axis=1)，示例代码与输出如下：

import polars as pl

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

# 获取每行水平最大值
result_max = df.with_columns(max = pl.max_horizontal("a", "b"))
print(result_max)

输出：

┌─────┬──────┬─────┐
│ a   ┆ b    ┆ max │
│ --- ┆ ---  ┆ --- │
│ i64 ┆ i64  ┆ i64 │
╞═════╪══════╪═════╡
│ 1   ┆ 4    ┆ 4   │
│ 8   ┆ 5    ┆ 8   │
│ 3   ┆ null ┆ 3   │
└─────┴──────┴─────┘

如果需要获取最大值对应的列名（对应Pandas的df[["a", "b"]].idxmax(axis=1)），可以直接使用Polars的pl.arg_max_horizontal函数，它会返回每行最大值所在的列名，自动忽略null值：

# 获取每行最大值对应的列名
result_idxmax = df.with_columns(max_col = pl.arg_max_horizontal("a", "b"))
print(result_idxmax)

预期输出：

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

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