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

Closed
Show file tree
Hide file tree
Changes from all commits
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 @@ -8,6 +8,8 @@
- alias
- all
- any
- arg_min
- arg_max
- arg_true
- cast
- count
Expand Down
2 changes: 2 additions & 0 deletions docs/api-reference/narwhals.md
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,8 @@ Here are the top-level functions available in Narwhals.
- all
- all_horizontal
- any_horizontal
- arg_max
- arg_min
- col
- concat
- concat_str
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 @@ -11,6 +11,8 @@
- alias
- all
- any
- arg_min
- arg_max
- arg_true
- cast
- clip
Expand Down
4 changes: 4 additions & 0 deletions narwhals/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,8 @@
from narwhals.expr import all_ as all
from narwhals.expr import all_horizontal
from narwhals.expr import any_horizontal
from narwhals.expr import arg_max
from narwhals.expr import arg_min
from narwhals.expr import col
from narwhals.expr import concat_str
from narwhals.expr import len_ as len
Expand Down Expand Up @@ -105,6 +107,8 @@
"all",
"all_horizontal",
"any_horizontal",
"arg_max",
"arg_min",
"col",
"concat",
"concat_str",
Expand Down
194 changes: 194 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).arg_min())

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).arg_max())

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

Expand Down Expand Up @@ -5865,6 +5957,57 @@ def median(*columns: str) -> Expr:
return Expr(lambda plx: plx.median(*columns))


def arg_min(*columns: str) -> Expr:
"""Return the index of the minimum value.

Note:
Syntactic sugar for ``nw.col(columns).argmin()``.

Arguments:
columns: Name(s) of the columns to use in the aggregation function.

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, 2], "b": [5, 10]})
>>> df_pl = pl.DataFrame({"a": [1, 2], "b": [5, 10]})
>>> df_pa = pa.table({"a": [1, 2], "b": [5, 10]})

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.argmin("b")).to_native()

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

>>> my_library_agnostic_function(df_pd)
b
0 0
>>> my_library_agnostic_function(df_pl)
shape: (1, 1)
β”Œβ”€β”€β”€β”€β”€β”
β”‚ b β”‚
β”‚ --- β”‚
β”‚ u32 β”‚
β•žβ•β•β•β•β•β•‘
β”‚ 0 β”‚
β””β”€β”€β”€β”€β”€β”˜
>>> my_library_agnostic_function(df_pa)
pyarrow.Table
b: int32
----
b: [[0]]
"""
return Expr(lambda plx: plx.arg_min(*columns))


def min(*columns: str) -> Expr:
"""Return the minimum value.

Expand Down Expand Up @@ -5916,6 +6059,57 @@ def min(*columns: str) -> Expr:
return Expr(lambda plx: plx.min(*columns))


def arg_max(*columns: str) -> Expr:
"""Return the index of the maximum value.

Note:
Syntactic sugar for ``nw.col(columns).argmax()``.

Arguments:
columns: Name(s) of the columns to use in the aggregation function.

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, 3, 2], "b": [5, 10, 1]})
>>> df_pl = pl.DataFrame({"a": [1, 3, 2], "b": [5, 10, 1]})
>>> df_pa = pa.table({"a": [1, 3, 2], "b": [5, 10, 1]})

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.argmax("a")).to_native()

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

>>> my_library_agnostic_function(df_pd)
a
0 1
>>> my_library_agnostic_function(df_pl)
shape: (1, 1)
β”Œβ”€β”€β”€β”€β”€β”
β”‚ a β”‚
β”‚ --- β”‚
β”‚ u32 β”‚
β•žβ•β•β•β•β•β•‘
β”‚ 1 β”‚
β””β”€β”€β”€β”€β”€β”˜
>>> my_library_agnostic_function(df_pa)
pyarrow.Table
a: int64
----
a: [[1]]
"""
return Expr(lambda plx: plx.arg_max(*columns))


def max(*columns: str) -> Expr:
"""Return the maximum value.

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.argmin()

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._from_compliant_series(self._compliant_series.arg_min())

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._from_compliant_series(self._compliant_series.arg_max())

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

Expand Down
Loading
Loading