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

Description

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

Usage

<Expr>$arr$arg_max()

Value

A polars expression

Examples

library("polars")

df <- pl$DataFrame(
  values = list(1:2, 2:1)
)$cast(pl$Array(pl$Int32, 2))
df$with_columns(
  arg_max = pl$col("values")$arr$arg_max()
)

#> shape: (2, 2)
#> ┌───────────────┬─────────┐
#> │ values        ┆ arg_max │
#> │ ---           ┆ ---     │
#> │ array[i32, 2] ┆ u32     │
#> ╞═══════════════╪═════════╡
#> │ [1, 2]        ┆ 1       │
#> │ [2, 1]        ┆ 0       │
#> └───────────────┴─────────┘