feat: add Series and Expr arg_min & arg_max (#1529)

---------

Co-authored-by: devdanzin <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Marco Gorelli <[email protected]>

Author: 4 people

docs/api-reference/expr.md

@@ -8,6 +8,8 @@
         - alias
         - all
         - any
+        - arg_max
+        - arg_min
         - arg_true
         - cast
         - count

docs/api-reference/series.md

@@ -11,6 +11,8 @@
         - alias
         - all
         - any
+        - arg_max
+        - arg_min
         - arg_true
         - cast
         - clip

narwhals/_arrow/expr.py

@@ -272,6 +272,12 @@ def min(self: Self) -> Self:
     def max(self: Self) -> Self:
         return reuse_series_implementation(self, "max", returns_scalar=True)
 
+    def arg_min(self: Self) -> Self:
+        return reuse_series_implementation(self, "arg_min", returns_scalar=True)
+
+    def arg_max(self: Self) -> Self:
+        return reuse_series_implementation(self, "arg_max", returns_scalar=True)
+
     def all(self: Self) -> Self:
         return reuse_series_implementation(self, "all", returns_scalar=True)
 

narwhals/_arrow/series.py

@@ -288,6 +288,18 @@ def max(self: Self, *, _return_py_scalar: bool = True) -> int:
 
         return maybe_extract_py_scalar(pc.max(self._native_series), _return_py_scalar)  # type: ignore[no-any-return]
 
+    def arg_min(self: Self, *, _return_py_scalar: bool = True) -> int:
+        import pyarrow.compute as pc
+
+        index_min = pc.index(self._native_series, pc.min(self._native_series))
+        return maybe_extract_py_scalar(index_min, _return_py_scalar)  # type: ignore[no-any-return]
+
+    def arg_max(self: Self, *, _return_py_scalar: bool = True) -> int:
+        import pyarrow.compute as pc
+
+        index_max = pc.index(self._native_series, pc.max(self._native_series))
+        return maybe_extract_py_scalar(index_max, _return_py_scalar)  # type: ignore[no-any-return]
+
     def sum(self: Self, *, _return_py_scalar: bool = True) -> int:
         import pyarrow.compute as pc
 

narwhals/_pandas_like/expr.py

@@ -264,6 +264,12 @@ def max(self) -> Self:
     def min(self) -> Self:
         return reuse_series_implementation(self, "min", returns_scalar=True)
 
+    def arg_min(self) -> Self:
+        return reuse_series_implementation(self, "arg_min", returns_scalar=True)
+
+    def arg_max(self) -> Self:
+        return reuse_series_implementation(self, "arg_max", returns_scalar=True)
+
     # Other
 
     def clip(self, lower_bound: Any, upper_bound: Any) -> Self:

narwhals/_pandas_like/series.py

@@ -284,6 +284,18 @@ def arg_true(self) -> PandasLikeSeries:
         result = ser.__class__(range(len(ser)), name=ser.name, index=ser.index).loc[ser]
         return self._from_native_series(result)
 
+    def arg_min(self) -> int:
+        ser = self._native_series
+        if self._implementation is Implementation.PANDAS and self._backend_version < (1,):
+            return ser.values.argmin()  # type: ignore[no-any-return]  # noqa: PD011
+        return ser.argmin()  # type: ignore[no-any-return]
+
+    def arg_max(self) -> int:
+        ser = self._native_series
+        if self._implementation is Implementation.PANDAS and self._backend_version < (1,):
+            return ser.values.argmax()  # type: ignore[no-any-return]  # noqa: PD011
+        return ser.argmax()  # type: ignore[no-any-return]
+
     # Binary comparisons
 
     def filter(self, other: Any) -> PandasLikeSeries:

narwhals/expr.py

@@ -962,6 +962,102 @@ def max(self) -> Self:
         """
         return self.__class__(lambda plx: self._to_compliant_expr(plx).max())
 
+    def arg_min(self) -> Self:
+        """Returns the index of the minimum value.
+
+        Returns:
+            A new expression.
+
+        Examples:
+            >>> import polars as pl
+            >>> import pandas as pd
+            >>> import pyarrow as pa
+            >>> import narwhals as nw
+            >>> from narwhals.typing import IntoFrameT
+            >>> df_pd = pd.DataFrame({"a": [10, 20], "b": [150, 100]})
+            >>> df_pl = pl.DataFrame({"a": [10, 20], "b": [150, 100]})
+            >>> df_pa = pa.table({"a": [10, 20], "b": [150, 100]})
+
+            Let's define a dataframe-agnostic function:
+
+            >>> def agnostic_arg_min(df_native: IntoFrameT) -> IntoFrameT:
+            ...     df = nw.from_native(df_native)
+            ...     return df.select(
+            ...         nw.col("a", "b").arg_min().name.suffix("_arg_min")
+            ...     ).to_native()
+
+            We can then pass any supported library such as Pandas, Polars, or PyArrow:
+
+            >>> agnostic_arg_min(df_pd)
+               a_arg_min  b_arg_min
+            0          0          1
+            >>> agnostic_arg_min(df_pl)
+            shape: (1, 2)
+            ┌───────────┬───────────┐
+            │ a_arg_min ┆ b_arg_min │
+            │ ---       ┆ ---       │
+            │ u32       ┆ u32       │
+            ╞═══════════╪═══════════╡
+            │ 0         ┆ 1         │
+            └───────────┴───────────┘
+            >>> agnostic_arg_min(df_pa)
+            pyarrow.Table
+            a_arg_min: int64
+            b_arg_min: int64
+            ----
+            a_arg_min: [[0]]
+            b_arg_min: [[1]]
+        """
+        return self.__class__(lambda plx: self._to_compliant_expr(plx).arg_min())
+
+    def arg_max(self) -> Self:
+        """Returns the index of the maximum value.
+
+        Returns:
+            A new expression.
+
+        Examples:
+            >>> import polars as pl
+            >>> import pandas as pd
+            >>> import pyarrow as pa
+            >>> import narwhals as nw
+            >>> from narwhals.typing import IntoFrameT
+            >>> df_pd = pd.DataFrame({"a": [10, 20], "b": [150, 100]})
+            >>> df_pl = pl.DataFrame({"a": [10, 20], "b": [150, 100]})
+            >>> df_pa = pa.table({"a": [10, 20], "b": [150, 100]})
+
+            Let's define a dataframe-agnostic function:
+
+            >>> def agnostic_arg_max(df_native: IntoFrameT) -> IntoFrameT:
+            ...     df = nw.from_native(df_native)
+            ...     return df.select(
+            ...         nw.col("a", "b").arg_max().name.suffix("_arg_max")
+            ...     ).to_native()
+
+            We can then pass any supported library such as Pandas, Polars, or PyArrow:
+
+            >>> agnostic_arg_max(df_pd)
+               a_arg_max  b_arg_max
+            0          1          0
+            >>> agnostic_arg_max(df_pl)
+            shape: (1, 2)
+            ┌───────────┬───────────┐
+            │ a_arg_max ┆ b_arg_max │
+            │ ---       ┆ ---       │
+            │ u32       ┆ u32       │
+            ╞═══════════╪═══════════╡
+            │ 1         ┆ 0         │
+            └───────────┴───────────┘
+            >>> agnostic_arg_max(df_pa)
+            pyarrow.Table
+            a_arg_max: int64
+            b_arg_max: int64
+            ----
+            a_arg_max: [[1]]
+            b_arg_max: [[0]]
+        """
+        return self.__class__(lambda plx: self._to_compliant_expr(plx).arg_max())
+
     def count(self) -> Self:
         """Returns the number of non-null elements in the column.
 

narwhals/series.py

@@ -931,6 +931,70 @@ def max(self) -> Any:
         """
         return self._compliant_series.max()
 
+    def arg_min(self) -> int:
+        """Returns the index of the minimum value.
+
+        Examples:
+            >>> import pandas as pd
+            >>> import polars as pl
+            >>> import pyarrow as pa
+            >>> import narwhals as nw
+            >>> from narwhals.typing import IntoSeries
+            >>> s = [1, 2, 3]
+            >>> s_pd = pd.Series(s)
+            >>> s_pl = pl.Series(s)
+            >>> s_pa = pa.chunked_array([s])
+
+            We define a library agnostic function:
+
+            >>> def agnostic_arg_min(s_native: IntoSeries):
+            ...     s = nw.from_native(s_native, series_only=True)
+            ...     return s.arg_min()
+
+            We can then pass either any supported library such as pandas, Polars,
+            or PyArrow:
+
+            >>> agnostic_arg_min(s_pd)
+            np.int64(0)
+            >>> agnostic_arg_min(s_pl)
+            0
+            >>> agnostic_arg_min(s_pa)
+            0
+        """
+        return self._compliant_series.arg_min()  # type: ignore[no-any-return]
+
+    def arg_max(self) -> int:
+        """Returns the index of the maximum value.
+
+        Examples:
+            >>> import pandas as pd
+            >>> import polars as pl
+            >>> import pyarrow as pa
+            >>> import narwhals as nw
+            >>> from narwhals.typing import IntoSeries
+            >>> s = [1, 2, 3]
+            >>> s_pd = pd.Series(s)
+            >>> s_pl = pl.Series(s)
+            >>> s_pa = pa.chunked_array([s])
+
+            We define a library agnostic function:
+
+            >>> def agnostic_arg_max(s_native: IntoSeries):
+            ...     s = nw.from_native(s_native, series_only=True)
+            ...     return s.arg_max()
+
+            We can then pass either any supported library such as pandas, Polars,
+            or PyArrow:
+
+            >>> agnostic_arg_max(s_pd)
+            np.int64(2)
+            >>> agnostic_arg_max(s_pl)
+            2
+            >>> agnostic_arg_max(s_pa)
+            2
+        """
+        return self._compliant_series.arg_max()  # type: ignore[no-any-return]
+
     def sum(self) -> Any:
         """Reduce this Series to the sum value.
 

pyproject.toml

@@ -142,20 +142,19 @@ omit = [
   'narwhals/typing.py',
   'narwhals/stable/v1/typing.py',
   'narwhals/this.py',
-  # we can run this in every environment that we measure coverage on due to upper-bound constraits
+  # we can't run this in every environment that we measure coverage on due to upper-bound constraits
   'narwhals/_ibis/*',
   # the latest pyspark (3.5) doesn't officially support Python 3.12 and 3.13
   'narwhals/_spark_like/*',
   'tests/spark_like_test.py',
 ]
 exclude_also = [
-  "> POLARS_VERSION",
   "if sys.version_info() <",
   "if (:?self._)?implementation is Implementation.MODIN",
   "if (:?self._)?implementation is Implementation.CUDF",
   'request.applymarker\(pytest.mark.xfail\)',
-  'if \w+._backend_version < ',
-  'if backend_version <',
+  '\w+._backend_version < ',
+  'backend_version <',
   'if "cudf" in str\(constructor'
 ]
 

tests/expr_and_series/arg_max_test.py

@@ -0,0 +1,41 @@
+from __future__ import annotations
+
+import pytest
+
+import narwhals.stable.v1 as nw
+from tests.utils import Constructor
+from tests.utils import ConstructorEager
+from tests.utils import assert_equal_data
+
+data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9], "i": [3, 1, 5]}
+
+
+def test_expr_arg_max_expr(
+    constructor: Constructor, request: pytest.FixtureRequest
+) -> None:
+    if "dask" in str(constructor):
+        # This operation is row-order dependent so we don't support it for Dask
+        request.applymarker(pytest.mark.xfail)
+    if "modin" in str(constructor):
+        # TODO(unassigned): bug in modin?
+        return
+    df = nw.from_native(constructor(data))
+    df = nw.maybe_set_index(df, "i")
+    result = df.select(nw.col("a", "b", "z").arg_max())
+    expected = {"a": [1], "b": [2], "z": [2]}
+    assert_equal_data(result, expected)
+
+
+@pytest.mark.parametrize(("col", "expected"), [("a", 1), ("b", 2), ("z", 2)])
+def test_expr_arg_max_series(
+    constructor_eager: ConstructorEager,
+    col: str,
+    expected: float,
+) -> None:
+    if "modin" in str(constructor_eager):
+        # TODO(unassigned): bug in modin?
+        return
+    series = nw.from_native(constructor_eager(data), eager_only=True)[col]
+    series = nw.maybe_set_index(series, index=[1, 0, 9])  # type: ignore[arg-type]
+    result = series.arg_max()
+    assert_equal_data({col: [result]}, {col: [expected]})

tests/expr_and_series/arg_min_test.py

@@ -0,0 +1,37 @@
+from __future__ import annotations
+
+import pytest
+
+import narwhals.stable.v1 as nw
+from tests.utils import Constructor
+from tests.utils import ConstructorEager
+from tests.utils import assert_equal_data
+
+data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]}
+
+
+def test_expr_arg_min_expr(
+    constructor: Constructor, request: pytest.FixtureRequest
+) -> None:
+    if "dask" in str(constructor):
+        # This operation is row-order dependent so we don't support it for Dask
+        request.applymarker(pytest.mark.xfail)
+    df = nw.from_native(constructor(data))
+    result = df.select(nw.col("a", "b", "z").arg_min())
+    expected = {"a": [0], "b": [0], "z": [0]}
+    assert_equal_data(result, expected)
+
+
+@pytest.mark.parametrize(("col", "expected"), [("a", 0), ("b", 0), ("z", 0)])
+def test_expr_arg_min_series(
+    constructor_eager: ConstructorEager,
+    col: str,
+    expected: float,
+) -> None:
+    if "modin" in str(constructor_eager):
+        # TODO(unassigned): bug in modin?
+        return
+    series = nw.from_native(constructor_eager(data), eager_only=True)[col]
+    series = nw.maybe_set_index(series, index=[1, 0, 9])  # type: ignore[arg-type]
+    result = series.arg_min()
+    assert_equal_data({col: [result]}, {col: [expected]})