Skip to content

Commit

Permalink
add invert to expressions, add tests for invert and sample
Browse files Browse the repository at this point in the history
  • Loading branch information
raisa committed Mar 20, 2024
1 parent ce542f2 commit 998ebe5
Show file tree
Hide file tree
Showing 2 changed files with 23 additions and 0 deletions.
3 changes: 3 additions & 0 deletions narwhals/expression.py
Original file line number Diff line number Diff line change
Expand Up @@ -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())

Expand Down
20 changes: 20 additions & 0 deletions tests/test_common.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)

0 comments on commit 998ebe5

Please sign in to comment.