Skip to content

Commit

Permalink
feat: Series.pipe and Expr.pipe (#582)
Browse files Browse the repository at this point in the history
  • Loading branch information
@muddi900
muddi900 authored Aug 9, 2024
1 parent ffef959 commit 5334934
Show file tree
Hide file tree
Showing 6 changed files with 166 additions and 61 deletions.
1 change: 1 addition & 0 deletions docs/api-reference/expr.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@
- null_count
- n_unique
- over
- pipe
- quantile
- round
- sample
Expand Down
1 change: 1 addition & 0 deletions docs/api-reference/series.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@
- name
- null_count
- n_unique
- pipe
- quantile
- round
- sample
Expand Down
8 changes: 4 additions & 4 deletions narwhals/dataframe.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1999,8 +1999,8 @@ def gather_every(self: Self, n: int, offset: int = 0) -> Self:
starting from a offset of 1:
>>> @nw.narwhalify
... def func(df_any):
... return df_any.gather_every(n=2, offset=1)
... def func(df):
... return df.gather_every(n=2, offset=1)
>>> func(df_pd)
a b
Expand Down Expand Up @@ -3299,8 +3299,8 @@ def gather_every(self: Self, n: int, offset: int = 0) -> Self:
starting from a offset of 1:
>>> @nw.narwhalify
... def func(df_any):
... return df_any.gather_every(n=2, offset=1)
... def func(df):
... return df.gather_every(n=2, offset=1)
>>> func(df_pd)
a b
Expand Down
154 changes: 99 additions & 55 deletions narwhals/expr.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,47 @@ def alias(self, name: str) -> Self:
"""
return self.__class__(lambda plx: self._call(plx).alias(name))

def pipe(self, function: Callable[[Any], Self], *args: Any, **kwargs: Any) -> Self:
"""
Pipe function call.
Examples:
>>> import polars as pl
>>> import pandas as pd
>>> import narwhals as nw
>>> data = {"a": [1, 2, 3, 4]}
>>> df_pd = pd.DataFrame(data)
>>> df_pl = pl.DataFrame(data)
Lets define a library-agnostic function:
>>> @nw.narwhalify
... def func(df):
... return df.select(nw.col("a").pipe(lambda x: x + 1))
We can then pass any supported library:
>>> func(df_pd)
a
0 2
1 3
2 4
3 5
>>> func(df_pl)
shape: (4, 1)
┌─────┐
│ a │
│ --- │
│ i64 │
╞═════╡
│ 2 │
│ 3 │
│ 4 │
│ 5 │
└─────┘
"""
return function(self, *args, **kwargs)

def cast(
self,
dtype: Any,
Expand Down Expand Up @@ -1309,13 +1350,13 @@ def is_duplicated(self) -> Self:
We can then pass either pandas or Polars to `func`:
>>> func(df_pd) # doctest: +NORMALIZE_WHITESPACE
>>> func(df_pd)
a b
0 True True
1 False True
2 False False
3 True False
>>> func(df_pl) # doctest: +NORMALIZE_WHITESPACE
>>> func(df_pl)
shape: (4, 2)
┌───────┬───────┐
│ a ┆ b │
Expand Down Expand Up @@ -1350,13 +1391,13 @@ def is_unique(self) -> Self:
We can then pass either pandas or Polars to `func`:
>>> func(df_pd) # doctest: +NORMALIZE_WHITESPACE
>>> func(df_pd)
a b
0 False False
1 True False
2 True True
3 False True
>>> func(df_pl) # doctest: +NORMALIZE_WHITESPACE
>>> func(df_pl)
shape: (4, 2)
┌───────┬───────┐
│ a ┆ b │
Expand Down Expand Up @@ -1431,13 +1472,13 @@ def is_first_distinct(self) -> Self:
We can then pass either pandas or Polars to `func`:
>>> func(df_pd) # doctest: +NORMALIZE_WHITESPACE
>>> func(df_pd)
a b
0 True True
1 True False
2 True True
3 False True
>>> func(df_pl) # doctest: +NORMALIZE_WHITESPACE
>>> func(df_pl)
shape: (4, 2)
┌───────┬───────┐
│ a ┆ b │
Expand Down Expand Up @@ -1471,13 +1512,13 @@ def is_last_distinct(self) -> Self:
We can then pass either pandas or Polars to `func`:
>>> func(df_pd) # doctest: +NORMALIZE_WHITESPACE
>>> func(df_pd)
a b
0 False False
1 True True
2 True True
3 True True
>>> func(df_pl) # doctest: +NORMALIZE_WHITESPACE
>>> func(df_pl)
shape: (4, 2)
┌───────┬───────┐
│ a ┆ b │
Expand Down Expand Up @@ -1524,11 +1565,11 @@ def quantile(
We can then pass either pandas or Polars to `func`:
>>> func(df_pd) # doctest: +NORMALIZE_WHITESPACE
a b
>>> func(df_pd)
a b
0 24.5 74.5
>>> func(df_pl) # doctest: +NORMALIZE_WHITESPACE
>>> func(df_pl)
shape: (1, 2)
┌──────┬──────┐
│ a ┆ b │
Expand Down Expand Up @@ -1566,12 +1607,12 @@ def head(self, n: int = 10) -> Self:
We can then pass either pandas or Polars to `func`:
>>> func(df_pd) # doctest: +NORMALIZE_WHITESPACE
>>> func(df_pd)
a
0 0
1 1
2 2
>>> func(df_pl) # doctest: +NORMALIZE_WHITESPACE
>>> func(df_pl)
shape: (3, 1)
┌─────┐
│ a │
Expand Down Expand Up @@ -1610,12 +1651,12 @@ def tail(self, n: int = 10) -> Self:
We can then pass either pandas or Polars to `func`:
>>> func(df_pd) # doctest: +NORMALIZE_WHITESPACE
a
>>> func(df_pd)
a
7 7
8 8
9 9
>>> func(df_pl) # doctest: +NORMALIZE_WHITESPACE
>>> func(df_pl)
shape: (3, 1)
┌─────┐
│ a │
Expand Down Expand Up @@ -1662,12 +1703,12 @@ def round(self, decimals: int = 0) -> Self:
We can then pass either pandas or Polars to `func`:
>>> func(df_pd) # doctest: +NORMALIZE_WHITESPACE
>>> func(df_pd)
a
0 1.1
1 2.6
2 3.9
>>> func(df_pl) # doctest: +NORMALIZE_WHITESPACE
>>> func(df_pl)
shape: (3, 1)
┌─────┐
│ a │
Expand Down Expand Up @@ -1707,10 +1748,10 @@ def len(self) -> Self:
We can then pass either pandas or Polars to `func`:
>>> func(df_pd) # doctest: +NORMALIZE_WHITESPACE
a1 a2
0 2 1
>>> func(df_pl) # doctest: +NORMALIZE_WHITESPACE
>>> func(df_pd)
a1 a2
0 2 1
>>> func(df_pl)
shape: (1, 2)
┌─────┬─────┐
│ a1 ┆ a2 │
Expand Down Expand Up @@ -1801,7 +1842,7 @@ def clip(
0 2
1 2
2 3
>>> func_lower(df_pl) # doctest: +NORMALIZE_WHITESPACE
>>> func_lower(df_pl)
shape: (3, 1)
┌─────┐
│ s │
Expand All @@ -1826,7 +1867,7 @@ def clip(
0 1
1 2
2 2
>>> func_upper(df_pl) # doctest: +NORMALIZE_WHITESPACE
>>> func_upper(df_pl)
shape: (3, 1)
┌─────┐
│ s │
Expand Down Expand Up @@ -1860,19 +1901,19 @@ def clip(
3 3
4 -1
5 3
>>> func(df_pl) # doctest: +NORMALIZE_WHITESPACE
>>> func(df_pl)
shape: (6, 1)
┌─────┐
│ s │
│ --- │
│ i64 │
╞═════╡
│ -1 │
│ 1 │
│ 1 │
│ -1 │
│ 3 │
│ 3 │
│ -1 │
│ 3 │
│ 3 │
└─────┘
"""
return self.__class__(
Expand Down Expand Up @@ -2162,7 +2203,7 @@ def slice(self, offset: int, length: int | None = None) -> Expr:
2 papaya ya
3 dragonfruit onf
>>> func(df_pl) # doctest: +NORMALIZE_WHITESPACE
>>> func(df_pl)
shape: (4, 2)
┌─────────────┬──────────┐
│ s ┆ s_sliced │
Expand All @@ -2181,14 +2222,14 @@ def slice(self, offset: int, length: int | None = None) -> Expr:
... def func(df):
... return df.with_columns(s_sliced=nw.col("s").str.slice(-3))
>>> func(df_pd) # doctest: +NORMALIZE_WHITESPACE
>>> func(df_pd)
s s_sliced
0 pear ear
1 None None
2 papaya aya
3 dragonfruit uit
>>> func(df_pl) # doctest: +NORMALIZE_WHITESPACE
>>> func(df_pl)
shape: (4, 2)
┌─────────────┬──────────┐
│ s ┆ s_sliced │
Expand Down Expand Up @@ -2375,8 +2416,8 @@ def to_uppercase(self) -> Expr:
We can then pass either pandas or Polars to `func`:
>>> func(df_pd) # doctest: +NORMALIZE_WHITESPACE
fruits upper_col
>>> func(df_pd)
fruits upper_col
0 apple APPLE
1 mango MANGO
2 None None
Expand Down Expand Up @@ -2416,13 +2457,13 @@ def to_lowercase(self) -> Expr:
We can then pass either pandas or Polars to `func`:
>>> func(df_pd) # doctest: +NORMALIZE_WHITESPACE
fruits lower_col
0 APPLE apple
1 MANGO mango
2 None None
>>> func(df_pd)
fruits lower_col
0 APPLE apple
1 MANGO mango
2 None None
>>> func(df_pl) # doctest: +NORMALIZE_WHITESPACE
>>> func(df_pl)
shape: (3, 2)
┌────────┬───────────┐
│ fruits ┆ lower_col │
Expand Down Expand Up @@ -2453,32 +2494,35 @@ def date(self) -> Expr:
>>> import polars as pl
>>> from datetime import datetime
>>> import narwhals as nw
>>> dates = [datetime(2012, 1, 7, 10, 20), datetime(2023, 3, 10, 11, 32)]
>>> s_pd = pd.Series(dates).convert_dtypes(
>>> data = {"a": [datetime(2012, 1, 7, 10, 20), datetime(2023, 3, 10, 11, 32)]}
>>> df_pd = pd.DataFrame(data).convert_dtypes(
... dtype_backend="pyarrow"
... ) # doctest:+SKIP
>>> s_pl = pl.Series(dates)
>>> df_pl = pl.DataFrame(data)
We define a library agnostic function:
>>> @nw.narwhalify
... def func(s):
... return s.dt.date()
... def func(df):
... return df.select(nw.col("a").dt.date())
We can then pass either pandas or Polars to `func`:
>>> func(s_pd) # doctest:+SKIP
0 2012-01-07
1 2023-03-10
dtype: date32[day][pyarrow]
>>> func(df_pd) # doctest:+SKIP
a
0 2012-01-07
1 2023-03-10
>>> func(s_pl) # doctest: +NORMALIZE_WHITESPACE
shape: (2,)
Series: '' [date]
[
2012-01-07
2023-03-10
]
>>> func(df_pl) # docetst
shape: (2, 1)
┌────────────┐
│ a │
│ --- │
│ date │
╞════════════╡
│ 2012-01-07 │
│ 2023-03-10 │
└────────────┘
"""
return self._expr.__class__(lambda plx: self._expr._call(plx).dt.date())

Expand Down
Loading

0 comments on commit 5334934

Please sign in to comment.