From 998ebe5a1f7534857e2c8b0583599577b5d9c700 Mon Sep 17 00:00:00 2001
From: raisa <>
Date: Wed, 20 Mar 2024 10:31:52 +0000
Subject: [PATCH] add invert to expressions, add tests for invert and sample

---
 narwhals/expression.py |  3 +++
 tests/test_common.py   | 20 ++++++++++++++++++++
 2 files changed, 23 insertions(+)

diff --git a/narwhals/expression.py b/narwhals/expression.py
index bd9c2aee9..b786b2493 100644
--- a/narwhals/expression.py
+++ b/narwhals/expression.py
@@ -112,6 +112,9 @@ def __ge__(self, other: Any) -> Expr:
         )
 
     # --- unary ---
+    def __invert__(self) -> Expr:
+        return self.__class__(lambda plx: self._call(plx).__invert__())
+
     def mean(self) -> Expr:
         return self.__class__(lambda plx: self._call(plx).mean())
 
diff --git a/tests/test_common.py b/tests/test_common.py
index 78c160d6a..2337bcbe8 100644
--- a/tests/test_common.py
+++ b/tests/test_common.py
@@ -15,6 +15,8 @@
 df_pandas = pd.DataFrame({"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]})
 df_polars = pl.DataFrame({"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]})
 df_lazy = pl.LazyFrame({"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]})
+df_pandas_na = pd.DataFrame({"a": [None, 3, 2], "b": [4, 4, 6], "z": [7.0, None, 9]})
+df_lazy_na = pl.LazyFrame({"a": [None, 3, 2], "b": [4, 4, 6], "z": [7.0, None, 9]})
 
 if os.environ.get("CI", None):
     import modin.pandas as mpd
@@ -321,3 +323,21 @@ def test_expr_min_max(df_raw: Any) -> None:
     expected_max = {"a": [3], "b": [6], "z": [9]}
     compare_dicts(result_min, expected_min)
     compare_dicts(result_max, expected_max)
+
+
+@pytest.mark.parametrize("df_raw", [df_polars, df_pandas, df_mpd, df_lazy])
+def test_expr_sample(df_raw: Any) -> None:
+    df = nw.LazyFrame(df_raw)
+    result_shape = nw.to_native(df.select(nw.col("a", "b").sample(n=2))).collect().shape
+    expected = (2, 2)
+    assert result_shape == expected
+
+
+@pytest.mark.parametrize("df_raw", [df_pandas_na, df_lazy_na])
+def test_expr_na(df_raw: Any) -> None:
+    df = nw.LazyFrame(df_raw)
+    result_nna = nw.to_native(
+        df.filter((~nw.col("a").is_null()) & (~nw.col("z").is_null()))
+    )
+    expected = {"a": [2], "b": [6], "z": [9]}
+    compare_dicts(result_nna, expected)