Repository: narwhals-dev/narwhals

Issue:
feat: Series.arg_min, Series.arg_max, Expr.arg_min, Expr.arg_max
These would all be useful in Prophet, I think we should add them

Discussion:
Nathan-Bransby-NMT: @devdanzin I'm sorry, I'm not at home. 

I shouldn't have tried to merge from the mobile app 🤦‍♂️ Are you able to check out what happened? 

If not I will have a look when I get home.
devdanzin:  > I shouldn't have tried to merge from the mobile app 🤦‍♂️ Are you able to check out what happened?

Thanks for merging!

I don't see any issues with the merge of my PR, do you mean merging the other PR? I think it'll be OK, we (or the maintainers) can squash the commits.

I'll add the requested changes to the open PR in your fork.

MarcoGorelli: thanks al! we squash and merge PRs so no need to worry about about the individual commits
devdanzin: Thanks for the review and guidance!

I'm sending PRs to @Nathan-Bransby-NMT's fork so we can work together, you can take a look at https://github.com/Nathan-Bransby-NMT/narwhals/pull/3 if you want to preview the changes (currently addressing review comments and adding tests).
Nathan-Bransby-NMT: @MarcoGorelli Sorry I just got home, I'll work on this now
FBruzzesi: Hey @Nathan-Bransby-NMT @devdanzin I think this is extremely close to being approved and merge. I have a tiny request on the test side: I think it would be nice to make a custom index for pandas, so that we verify that it's not the index being picked up.

I would suggest to change the index via `nw.maybe_set_index`, e.g.:

```py
df = nw.maybe_set_index(nw.from_native(data), column_names=["a"]
```
and then exclude column `"a"` when performing `arg_max/min`
devdanzin: Thank you for the review and suggestion!

> I think it would be nice to make a custom index for pandas
> [...]
> I would suggest to change the index via `nw.maybe_set_index`, e.g.:
> [...]
> and then exclude column `"a"` when performing `arg_max/min`

I'm not sure about the best way of making a custom index for pandas only, does this make sense? Is this in the right ballpark of what you're asking for?

```python
def test_expr_arg_max_expr(constructor_eager: ConstructorEager) -> None:
    native = constructor_eager(data)
    if native.__class__ == PandasDataFrame:
        df = nw.maybe_set_index(nw.from_native(native), column_names=["a"])
    else:
        df = nw.from_native(native)
    result = df.select(nw.col("b", "z").arg_max())
    expected = {"b": [2], "z": [2]}
    assert_equal_data(result, expected)
```
MarcoGorelli: i think just
```
df = nw.maybe_set_index(nw.from_native(native), column_names=["a"])
```
in all cases probably works
MarcoGorelli: i think you could do it for series too, e.g.
```
s = nw.maybe_set_index(s, index=[1,2,3])
```

Patch:
diff --git a/docs/api-reference/expr.md b/docs/api-reference/expr.md
index 91f49a0a..61901a99 100644
--- a/docs/api-reference/expr.md
+++ b/docs/api-reference/expr.md
@@ -8,6 +8,8 @@
         - alias
         - all
         - any
+        - arg_max
+        - arg_min
         - arg_true
         - cast
         - count
diff --git a/docs/api-reference/series.md b/docs/api-reference/series.md
index 45c71f9c..79e7be8d 100644
--- a/docs/api-reference/series.md
+++ b/docs/api-reference/series.md
@@ -11,6 +11,8 @@
         - alias
         - all
         - any
+        - arg_max
+        - arg_min
         - arg_true
         - cast
         - clip
diff --git a/narwhals/_arrow/expr.py b/narwhals/_arrow/expr.py
index c875b7b5..ae8f2d03 100644
--- a/narwhals/_arrow/expr.py
+++ b/narwhals/_arrow/expr.py
@@ -272,6 +272,12 @@ class ArrowExpr(CompliantExpr[ArrowSeries]):
     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)
 
diff --git a/narwhals/_arrow/series.py b/narwhals/_arrow/series.py
index 8e3dacd6..900a6036 100644
--- a/narwhals/_arrow/series.py
+++ b/narwhals/_arrow/series.py
@@ -288,6 +288,18 @@ class ArrowSeries(CompliantSeries):
 
         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
 
diff --git a/narwhals/_pandas_like/expr.py b/narwhals/_pandas_like/expr.py
index 2f6ec16a..a83d83c4 100644
--- a/narwhals/_pandas_like/expr.py
+++ b/narwhals/_pandas_like/expr.py
@@ -264,6 +264,12 @@ class PandasLikeExpr(CompliantExpr[PandasLikeSeries]):
     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:
diff --git a/narwhals/_pandas_like/series.py b/narwhals/_pandas_like/series.py
index 9f5b562d..303583a9 100644
--- a/narwhals/_pandas_like/series.py
+++ b/narwhals/_pandas_like/series.py
@@ -284,6 +284,18 @@ class PandasLikeSeries(CompliantSeries):
         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:
diff --git a/narwhals/dataframe.py b/narwhals/dataframe.py
index 04cd43da..7ea45d59 100644
--- a/narwhals/dataframe.py
+++ b/narwhals/dataframe.py
@@ -536,6 +536,9 @@ class DataFrame(BaseFrame[DataFrameT]):
     def to_pandas(self) -> pd.DataFrame:
         """Convert this DataFrame to a pandas DataFrame.
 
+        Returns:
+            A pandas DataFrame.
+
         Examples:
             Construct pandas, Polars (eager) and PyArrow DataFrames:
 
@@ -581,8 +584,15 @@ class DataFrame(BaseFrame[DataFrameT]):
     def write_csv(self, file: str | Path | BytesIO | None = None) -> Any:
         r"""Write dataframe to comma-separated values (CSV) file.
 
+        Arguments:
+            file: String, path object or file-like object to which the dataframe will be
+                written. If None, the resulting csv format is returned as a string.
+
+        Returns:
+            String or None.
+
         Examples:
-            Construct pandas and Polars DataFrames:
+            Construct pandas, Polars (eager) and PyArrow DataFrames:
 
             >>> import pandas as pd
             >>> import polars as pl
@@ -613,9 +623,16 @@ class DataFrame(BaseFrame[DataFrameT]):
         """
         return self._compliant_frame.write_csv(file)
 
-    def write_parquet(self, file: str | Path | BytesIO) -> Any:
+    def write_parquet(self, file: str | Path | BytesIO) -> None:
         """Write dataframe to parquet file.
 
+        Arguments:
+            file: String, path object or file-like object to which the dataframe will be
+                written.
+
+        Returns:
+            None.
+
         Examples:
             Construct pandas, Polars and PyArrow DataFrames:
 
@@ -645,6 +662,9 @@ class DataFrame(BaseFrame[DataFrameT]):
     def to_numpy(self) -> np.ndarray:
         """Convert this DataFrame to a NumPy ndarray.
 
+        Returns:
+            A NumPy ndarray array.
+
         Examples:
             Construct pandas and polars DataFrames:
 
@@ -687,6 +707,9 @@ class DataFrame(BaseFrame[DataFrameT]):
     def shape(self) -> tuple[int, int]:
         """Get the shape of the DataFrame.
 
+        Returns:
+            The shape of the dataframe as a tuple.
+
         Examples:
             Construct pandas and polars DataFrames:
 
@@ -721,6 +744,12 @@ class DataFrame(BaseFrame[DataFrameT]):
     def get_column(self, name: str) -> Series[Any]:
         """Get a single column by name.
 
+        Arguments:
+            name: The column name as a string.
+
+        Returns:
+            A Narwhals Series, backed by a native series.
+
         Notes:
             Although `name` is typed as `str`, pandas does allow non-string column
             names, and they will work when passed to this function if the
@@ -883,6 +912,9 @@ class DataFrame(BaseFrame[DataFrameT]):
                     _inclusive_ and returns a `DataFrame`. For example, if the columns are
                     `'a', 'd', 'c', 'b'`, then that would extract columns `'a'`, `'d'`, and `'c'`.
 
+        Returns:
+            A Narwhals Series, backed by a native series.
+
         Notes:
             - Integers are always interpreted as positions
             - Strings are always interpreted as column names.
diff --git a/narwhals/expr.py b/narwhals/expr.py
index fc080470..674bb19a 100644
--- a/narwhals/expr.py
+++ b/narwhals/expr.py
@@ -962,6 +962,102 @@ class Expr:
         """
         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.
 
diff --git a/narwhals/series.py b/narwhals/series.py
index 602bb9bc..2797ceaf 100644
--- a/narwhals/series.py
+++ b/narwhals/series.py
@@ -931,6 +931,70 @@ class Series(Generic[IntoSeriesT]):
         """
         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.
 
diff --git a/pyproject.toml b/pyproject.toml
index 19330d76..bcf4671d 100644
--- a/pyproject.toml
+++ b/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'
 ]
 


Test patch:
diff --git a/tests/expr_and_series/arg_max_test.py b/tests/expr_and_series/arg_max_test.py
new file mode 100644
index 00000000..56e65aff
--- /dev/null
+++ b/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]})
diff --git a/tests/expr_and_series/arg_min_test.py b/tests/expr_and_series/arg_min_test.py
new file mode 100644
index 00000000..cde3b762
--- /dev/null
+++ b/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]})
diff --git a/tests/group_by_test.py b/tests/group_by_test.py
index 25c049ea..e51cbac6 100644
--- a/tests/group_by_test.py
+++ b/tests/group_by_test.py
@@ -98,24 +98,28 @@ def test_group_by_iter(constructor_eager: ConstructorEager) -> None:
 
 
 @pytest.mark.parametrize(
-    ("expr", "expected"),
+    ("attr", "expected"),
     [
-        (nw.col("b").sum(), {"a": [1, 2], "b": [3, 3]}),
-        (nw.col("b").mean(), {"a": [1, 2], "b": [1.5, 3]}),
-        (nw.col("b").max(), {"a": [1, 2], "b": [2, 3]}),
-        (nw.col("b").min(), {"a": [1, 2], "b": [1, 3]}),
-        (nw.col("b").std(), {"a": [1, 2], "b": [0.707107, None]}),
-        (nw.col("b").len(), {"a": [1, 2], "b": [3, 1]}),
-        (nw.col("b").n_unique(), {"a": [1, 2], "b": [3, 1]}),
-        (nw.col("b").count(), {"a": [1, 2], "b": [2, 1]}),
+        ("sum", {"a": [1, 2], "b": [3, 3]}),
+        ("mean", {"a": [1, 2], "b": [1.5, 3]}),
+        ("max", {"a": [1, 2], "b": [2, 3]}),
+        ("min", {"a": [1, 2], "b": [1, 3]}),
+        ("std", {"a": [1, 2], "b": [0.707107, None]}),
+        ("len", {"a": [1, 2], "b": [3, 1]}),
+        ("n_unique", {"a": [1, 2], "b": [3, 1]}),
+        ("count", {"a": [1, 2], "b": [2, 1]}),
     ],
 )
 def test_group_by_depth_1_agg(
     constructor: Constructor,
-    expr: nw.Expr,
+    attr: str,
     expected: dict[str, list[int | float]],
+    request: pytest.FixtureRequest,
 ) -> None:
+    if "cudf" in str(constructor) and attr == "n_unique":
+        request.applymarker(pytest.mark.xfail)
     data = {"a": [1, 1, 1, 2], "b": [1, None, 2, 3]}
+    expr = getattr(nw.col("b"), attr)()
     result = nw.from_native(constructor(data)).group_by("a").agg(expr).sort("a")
     assert_equal_data(result, expected)
 
@@ -346,6 +350,10 @@ def test_group_by_shift_raises(
         # Polars supports all kinds of crazy group-by aggregations, so
         # we don't check that it errors here.
         request.applymarker(pytest.mark.xfail)
+    if "cudf" in str(constructor):
+        # This operation fails completely in cuDF anyway, we just let raise its own
+        # error.
+        request.applymarker(pytest.mark.xfail)
     df_native = {"a": [1, 2, 3], "b": [1, 1, 2]}
     df = nw.from_native(constructor(df_native))
     with pytest.raises(