Repository: narwhals-dev/narwhals
Pull request #128 (2024-05-10 09:15:41)

Drop nulls and more


Patch:
diff --git a/docs/api-reference/dataframe.md b/docs/api-reference/dataframe.md
index 3f807456..40321868 100644
--- a/docs/api-reference/dataframe.md
+++ b/docs/api-reference/dataframe.md
@@ -6,6 +6,7 @@
       members:
         - columns
         - drop
+        - drop_nulls
         - filter
         - group_by
         - head
diff --git a/docs/api-reference/expressions.md b/docs/api-reference/expressions.md
index ba096720..aa4ab1e6 100644
--- a/docs/api-reference/expressions.md
+++ b/docs/api-reference/expressions.md
@@ -23,6 +23,7 @@
         - over
         - unique
         - sample
+        - shift
         - sort
         - std
         - sum
diff --git a/docs/api-reference/lazyframe.md b/docs/api-reference/lazyframe.md
index 7f858b34..d1ed8022 100644
--- a/docs/api-reference/lazyframe.md
+++ b/docs/api-reference/lazyframe.md
@@ -7,6 +7,7 @@
         - collect
         - columns
         - drop
+        - drop_nulls
         - filter
         - group_by
         - head
diff --git a/docs/api-reference/series.md b/docs/api-reference/series.md
index 84bcf363..fc35c1a9 100644
--- a/docs/api-reference/series.md
+++ b/docs/api-reference/series.md
@@ -24,6 +24,7 @@
         - n_unique
         - sample
         - shape
+        - shift
         - sort
         - std
         - sum
diff --git a/narwhals/_pandas_like/dataframe.py b/narwhals/_pandas_like/dataframe.py
index 983abb78..e97c9334 100644
--- a/narwhals/_pandas_like/dataframe.py
+++ b/narwhals/_pandas_like/dataframe.py
@@ -98,6 +98,9 @@ class PandasDataFrame:
         )
         return self._from_dataframe(df)
 
+    def drop_nulls(self) -> Self:
+        return self._from_dataframe(self._dataframe.dropna(axis=0))
+
     def filter(
         self,
         *predicates: IntoPandasExpr | Iterable[IntoPandasExpr],
diff --git a/narwhals/_pandas_like/expr.py b/narwhals/_pandas_like/expr.py
index 7c0c335e..cc7b6e8a 100644
--- a/narwhals/_pandas_like/expr.py
+++ b/narwhals/_pandas_like/expr.py
@@ -205,6 +205,9 @@ class PandasExpr:
     def diff(self) -> Self:
         return register_expression_call(self, "diff")
 
+    def shift(self, n: int) -> Self:
+        return register_expression_call(self, "shift", n)
+
     def sample(
         self,
         n: int | None = None,
diff --git a/narwhals/_pandas_like/series.py b/narwhals/_pandas_like/series.py
index 1dbac0e3..854ba161 100644
--- a/narwhals/_pandas_like/series.py
+++ b/narwhals/_pandas_like/series.py
@@ -375,6 +375,11 @@ class PandasSeries:
             self._series.__class__(self._series.diff(), name=self._series.name)
         )
 
+    def shift(self, n: int) -> PandasSeries:
+        return self._from_series(
+            self._series.__class__(self._series.shift(n), name=self._series.name)
+        )
+
     def sort(
         self,
         *,
diff --git a/narwhals/dataframe.py b/narwhals/dataframe.py
index 74736447..c3d08634 100644
--- a/narwhals/dataframe.py
+++ b/narwhals/dataframe.py
@@ -81,6 +81,11 @@ class BaseFrame:
     def pipe(self, function: Callable[[Any], Self], *args: Any, **kwargs: Any) -> Self:
         return function(self, *args, **kwargs)
 
+    def drop_nulls(self) -> Self:
+        return self._from_dataframe(
+            self._dataframe.drop_nulls(),
+        )
+
     @property
     def columns(self) -> list[str]:
         return self._dataframe.columns  # type: ignore[no-any-return]
@@ -477,6 +482,46 @@ class DataFrame(BaseFrame):
         """
         return super().pipe(function, *args, **kwargs)
 
+    def drop_nulls(self) -> Self:
+        """
+        Drop null values.
+
+        Notes:
+            pandas and Polars handle null values differently. Polars distinguishes
+            between NaN and Null, whereas pandas doesn't.
+
+        Examples:
+            >>> import polars as pl
+            >>> import pandas as pd
+            >>> import narwhals as nw
+            >>> data = {'a': [1., 2., None], 'ba': [1, None, 2.]}
+            >>> df_pd = pd.DataFrame(data)
+            >>> df_pl = pl.DataFrame(data)
+
+            Let's define a dataframe-agnostic function:
+
+            >>> def func(df_any):
+            ...     df = nw.from_native(df_any)
+            ...     df = df.drop_nulls()
+            ...     return nw.to_native(df)
+
+            We can then pass either pandas or Polars:
+
+            >>> func(df_pd)
+                 a   ba
+            0  1.0  1.0
+            >>> func(df_pl)
+            shape: (1, 2)
+            ┌─────┬─────┐
+            │ a   ┆ ba  │
+            │ --- ┆ --- │
+            │ f64 ┆ f64 │
+            ╞═════╪═════╡
+            │ 1.0 ┆ 1.0 │
+            └─────┴─────┘
+        """
+        return super().drop_nulls()
+
     @property
     def schema(self) -> dict[str, DType]:
         r"""
@@ -1426,6 +1471,46 @@ class LazyFrame(BaseFrame):
         """
         return super().pipe(function, *args, **kwargs)
 
+    def drop_nulls(self) -> Self:
+        """
+        Drop null values.
+
+        Notes:
+            pandas and Polars handle null values differently. Polars distinguishes
+            between NaN and Null, whereas pandas doesn't.
+
+        Examples:
+            >>> import polars as pl
+            >>> import pandas as pd
+            >>> import narwhals as nw
+            >>> data = {'a': [1., 2., None], 'ba': [1, None, 2.]}
+            >>> df_pd = pd.DataFrame(data)
+            >>> df_pl = pl.LazyFrame(data)
+
+            Let's define a dataframe-agnostic function:
+
+            >>> def func(df_any):
+            ...     df = nw.from_native(df_any)
+            ...     df = df.drop_nulls()
+            ...     return nw.to_native(df)
+
+            We can then pass either pandas or Polars:
+
+            >>> func(df_pd)
+                 a   ba
+            0  1.0  1.0
+            >>> func(df_pl).collect()
+            shape: (1, 2)
+            ┌─────┬─────┐
+            │ a   ┆ ba  │
+            │ --- ┆ --- │
+            │ f64 ┆ f64 │
+            ╞═════╪═════╡
+            │ 1.0 ┆ 1.0 │
+            └─────┴─────┘
+        """
+        return super().drop_nulls()
+
     @property
     def schema(self) -> dict[str, DType]:
         r"""
diff --git a/narwhals/expression.py b/narwhals/expression.py
index 65c8e3f6..cea122be 100644
--- a/narwhals/expression.py
+++ b/narwhals/expression.py
@@ -624,6 +624,60 @@ class Expr:
         """
         return self.__class__(lambda plx: self._call(plx).diff())
 
+    def shift(self, n: int) -> Expr:
+        """
+        Shift values by `n` positions.
+
+        Notes:
+            pandas may change the dtype here, for example when introducing missing
+            values in an integer column. To ensure, that the dtype doesn't change,
+            you may want to use `fill_null` and `cast`. For example, to shift
+            and fill missing values with `0` in a Int64 column, you could
+            do:
+
+            ```python
+            nw.col('a').shift(1).fill_null(0).cast(nw.Int64)
+            ```
+
+        Examples:
+            >>> import polars as pl
+            >>> import pandas as pd
+            >>> import narwhals as nw
+            >>> df_pd = pd.DataFrame({'a': [1, 1, 3, 5, 5]})
+            >>> df_pl = pl.DataFrame({'a': [1, 1, 3, 5, 5]})
+
+            Let's define a dataframe-agnostic function:
+
+            >>> def func(df_any):
+            ...    df = nw.from_native(df_any)
+            ...    df = df.select(a_shift=nw.col('a').shift(n=1))
+            ...    return nw.to_native(df)
+
+            We can then pass either pandas or Polars to `func`:
+
+            >>> func(df_pd)
+               a_shift
+            0      NaN
+            1      1.0
+            2      1.0
+            3      3.0
+            4      5.0
+            >>> func(df_pl)
+            shape: (5, 1)
+            ┌─────────┐
+            │ a_shift │
+            │ ---     │
+            │ i64     │
+            ╞═════════╡
+            │ null    │
+            │ 1       │
+            │ 1       │
+            │ 3       │
+            │ 5       │
+            └─────────┘
+        """
+        return self.__class__(lambda plx: self._call(plx).shift(n))
+
     def sort(self, *, descending: bool = False) -> Expr:
         return self.__class__(lambda plx: self._call(plx).sort(descending=descending))
 
diff --git a/narwhals/series.py b/narwhals/series.py
index ea843069..7f054b18 100644
--- a/narwhals/series.py
+++ b/narwhals/series.py
@@ -338,6 +338,54 @@ class Series:
         """
         return self._from_series(self._series.diff())
 
+    def shift(self, n: int) -> Self:
+        """
+        Shift values by `n` positions.
+
+        Notes:
+            pandas may change the dtype here, for example when introducing missing
+            values in an integer column. To ensure, that the dtype doesn't change,
+            you may want to use `fill_null` and `cast`. For example, to shift
+            and fill missing values with `0` in a Int64 column, you could
+            do:
+
+            ```python
+            s.shift(1).fill_null(0).cast(nw.Int64)
+            ```
+
+        Examples:
+            >>> import pandas as pd
+            >>> import polars as pl
+            >>> import narwhals as nw
+            >>> s = [2, 4, 3]
+            >>> s_pd = pd.Series(s)
+            >>> s_pl = pl.Series(s)
+
+            We define a dataframe-agnostic function:
+
+            >>> def func(s_any):
+            ...     s = nw.from_native(s_any, series_only=True)
+            ...     s = s.shift(1)
+            ...     return nw.to_native(s)
+
+            We can then pass either pandas or Polars to `func`:
+
+            >>> func(s_pd)
+            0    NaN
+            1    2.0
+            2    4.0
+            dtype: float64
+            >>> func(s_pl)  # doctest: +NORMALIZE_WHITESPACE
+            shape: (3,)
+            Series: '' [i64]
+            [
+               null
+               2
+               4
+            ]
+        """
+        return self._from_series(self._series.shift(n))
+
     def sample(
         self,
         n: int | None = None,


Test patch:
diff --git a/tests/expr/shift_test.py b/tests/expr/shift_test.py
new file mode 100644
index 00000000..05ad1402
--- /dev/null
+++ b/tests/expr/shift_test.py
@@ -0,0 +1,35 @@
+from typing import Any
+
+import pandas as pd
+import polars as pl
+import pytest
+
+import narwhals as nw
+from tests.utils import compare_dicts
+
+data = {
+    "i": [0, 1, 2, 3, 4],
+    "a": [0, 1, 2, 3, 4],
+    "b": [1, 2, 3, 5, 3],
+    "c": [5, 4, 3, 2, 1],
+}
+
+
+@pytest.mark.parametrize("constructor", [pd.DataFrame, pl.DataFrame])
+def test_over_single(constructor: Any) -> None:
+    df = nw.from_native(constructor(data), eager_only=True)
+    result = df.with_columns(nw.col("a", "b", "c").shift(2)).filter(nw.col("i") > 1)
+    expected = {
+        "i": [2, 3, 4],
+        "a": [0, 1, 2],
+        "b": [1, 2, 3],
+        "c": [5, 4, 3],
+    }
+    compare_dicts(result, expected)
+    result = df.select(
+        df["i"],
+        df["a"].shift(2),
+        df["b"].shift(2),
+        df["c"].shift(2),
+    ).filter(nw.col("i") > 1)
+    compare_dicts(result, expected)
diff --git a/tests/frame/drop_nulls_test.py b/tests/frame/drop_nulls_test.py
new file mode 100644
index 00000000..2a8a8dbd
--- /dev/null
+++ b/tests/frame/drop_nulls_test.py
@@ -0,0 +1,25 @@
+from typing import Any
+
+import pandas as pd
+import polars as pl
+import pytest
+
+import narwhals as nw
+from tests.utils import compare_dicts
+
+data = {
+    "a": [1.0, 2.0, None, 4.0],
+    "b": [None, 3.0, None, 5.0],
+}
+
+
+@pytest.mark.parametrize("constructor", [pd.DataFrame, pl.DataFrame])
+def test_drop_nulls(constructor: Any) -> None:
+    result = nw.from_native(constructor(data)).drop_nulls()
+    expected = {
+        "a": [2.0, 4.0],
+        "b": [3.0, 5.0],
+    }
+    compare_dicts(result, expected)
+    result = nw.from_native(constructor(data)).lazy().drop_nulls()
+    compare_dicts(result, expected)