polars.Expr.arr.arg_max#

Expr.arr.arg_max() → Expr[source]#

Retrieve an index of a maximum value in every sub-array.

When multiple values are equal to the maximum, this function may arbitrarily return the index of any of the maximum values. In this case, the returned index is not guaranteed to be the same across multiple runs.

engine:In MemoryStreamingDistributed

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       │
└───────────────┴─────────┘