Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Added the argmin and argmax function expr.py file and series.py file #1496

Open
wants to merge 6 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions docs/api-reference/expr.md
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,8 @@
- diff
- drop_nulls
- ewm_mean
- arg_min
- arg_max
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

these should be sorted alphabetically

- fill_null
- filter
- gather_every
Expand Down
2 changes: 2 additions & 0 deletions docs/api-reference/series.md
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,8 @@
- max
- mean
- median
- arg_min
- arg_max
- min
- mode
- name
Expand Down
92 changes: 92 additions & 0 deletions narwhals/expr.py
Original file line number Diff line number Diff line change
Expand Up @@ -621,6 +621,98 @@ def median(self) -> Self:
"""
return self.__class__(lambda plx: self._call(plx).median())

def arg_min(self) -> Self:
"""Get 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": [-1, 0, 1], "b": [2, 4, 6]})
>>> df_pl = pl.DataFrame({"a": [-1, 0, 1], "b": [2, 4, 6]})
>>> df_pa = pa.table({"a": [-1, 0, 1], "b": [2, 4, 6]})

Let's define a dataframe-agnostic function:

>>> def my_library_agnostic_function(df_native: IntoFrameT) -> IntoFrameT:
... df = nw.from_native(df_native)
... return df.select(nw.col("a", "b").argmin()).to_native()

We can then pass any supported library such as Pandas, Polars, or PyArrow to `func`:

>>> my_library_agnostic_function(df_pd)
a b
0 0 0
>>> my_library_agnostic_function(df_pl)
shape: (1, 2)
β”Œβ”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”
β”‚ a ┆ b β”‚
β”‚ --- ┆ --- β”‚
β”‚ i64 ┆ i64 β”‚
β•žβ•β•β•β•β•β•ͺ═════║
β”‚ 0 ┆ 0 β”‚
β””β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”˜
>>> my_library_agnostic_function(df_pa)
pyarrow.Table
a: int64
b: int64
----
a: [[0]]
b: [[0]]
"""
return self.__class__(lambda plx: self._call(plx).argmin())

def arg_max(self) -> Self:
"""Get 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": [-1, 0, 1], "b": [2, 4, 6]})
>>> df_pl = pl.DataFrame({"a": [-1, 0, 1], "b": [2, 4, 6]})
>>> df_pa = pa.table({"a": [-1, 0, 1], "b": [2, 4, 6]})

Let's define a dataframe-agnostic function:

>>> def my_library_agnostic_function(df_native: IntoFrameT) -> IntoFrameT:
... df = nw.from_native(df_native)
... return df.select(nw.col("a", "b").argmax()).to_native()

We can then pass any supported library such as Pandas, Polars, or PyArrow to `func`:

>>> my_library_agnostic_function(df_pd)
a b
2 2 2
>>> my_library_agnostic_function(df_pl)
shape: (1, 2)
β”Œβ”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”
β”‚ a ┆ b β”‚
β”‚ --- ┆ --- β”‚
β”‚ i64 ┆ i64 β”‚
β•žβ•β•β•β•β•β•ͺ═════║
β”‚ 2 ┆ 2 β”‚
β””β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”˜
>>> my_library_agnostic_function(df_pa)
pyarrow.Table
a: int64
b: int64
----
a: [[2]]
b: [[2]]
"""
return self.__class__(lambda plx: self._call(plx).argmax())

def std(self, *, ddof: int = 1) -> Self:
"""Get standard deviation.

Expand Down
54 changes: 54 additions & 0 deletions narwhals/series.py
Original file line number Diff line number Diff line change
Expand Up @@ -660,6 +660,60 @@ def median(self) -> Any:
"""
return self._compliant_series.median()

def arg_min(self) -> Any:
"""Get the index of the minimum value in this Series.

Examples:
>>> import pandas as pd
>>> import polars as pl
>>> import narwhals as nw
>>> from narwhals.typing import IntoSeries
>>> s = [1, 2, 3]
>>> s_pd = pd.Series(s)
>>> s_pl = pl.Series(s)

We define a library agnostic function:

>>> def my_library_agnostic_function(s_native: IntoSeries):
... s = nw.from_native(s_native, series_only=True)
... return s.arg_min()

We can then pass either pandas or Polars to `func`:

>>> my_library_agnostic_function(s_pd)
0
>>> my_library_agnostic_function(s_pl)
0
"""
return self._compliant_series.argmin()

def arg_max(self) -> Any:
"""Get the index of the maximum value in this Series.

Examples:
>>> import pandas as pd
>>> import polars as pl
>>> import narwhals as nw
>>> from narwhals.typing import IntoSeries
>>> s = [1, 2, 3]
>>> s_pd = pd.Series(s)
>>> s_pl = pl.Series(s)

We define a library agnostic function:

>>> def my_library_agnostic_function(s_native: IntoSeries):
... s = nw.from_native(s_native, series_only=True)
... return s.arg_max()

We can then pass either pandas or Polars to `func`:

>>> my_library_agnostic_function(s_pd)
2
>>> my_library_agnostic_function(s_pl)
2
"""
return self._compliant_series.argmax()

def skew(self: Self) -> Any:
"""Calculate the sample skewness of the Series.

Expand Down
Loading