polars.Expr.arr.arg_max

Expr.arr.arg_max() → Expr

Retrieve the index of the maximum value in every sub-array.

Returns:

Expr

Expression of data type UInt32 or UInt64 (depending on compilation).

Examples

>>> df = pl.DataFrame(
...     {
...         "a": [[1, 2], [2, 1]],
...     },
...     schema={"a": pl.Array(pl.Int64, 2)},
... )
>>> df.with_columns(arg_max=pl.col("a").arr.arg_max())
shape: (2, 2)
┌───────────────┬─────────┐
│ a             ┆ arg_max │
│ ---           ┆ ---     │
│ array[i64, 2] ┆ u32     │
╞═══════════════╪═════════╡
│ [1, 2]        ┆ 1       │
│ [2, 1]        ┆ 0       │
└───────────────┴─────────┘