Skip to content

Commit

Permalink
Added the argmin and argmax function expr.py file and series.py file
Browse files Browse the repository at this point in the history
  • Loading branch information
MUKESHRAJMAHENDRAN committed Dec 3, 2024
1 parent e02685b commit 7f939c6
Show file tree
Hide file tree
Showing 4 changed files with 150 additions and 0 deletions.
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
- 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

0 comments on commit 7f939c6

Please sign in to comment.