diff --git a/narwhals/expression.py b/narwhals/expression.py index 11a9cb0a5..bd9c2aee9 100644 --- a/narwhals/expression.py +++ b/narwhals/expression.py @@ -148,7 +148,13 @@ def is_null(self) -> Expr: def drop_nulls(self) -> Expr: return self.__class__(lambda plx: self._call(plx).drop_nulls()) - def sample(self, n: int, fraction: float, *, with_replacement: bool) -> Expr: + def sample( + self, + n: int | None = None, + fraction: float | None = None, + *, + with_replacement: bool = False, + ) -> Expr: return self.__class__( lambda plx: self._call(plx).sample( n, fraction=fraction, with_replacement=with_replacement diff --git a/tests/test_common.py b/tests/test_common.py index 0cef13264..c1034b738 100644 --- a/tests/test_common.py +++ b/tests/test_common.py @@ -257,7 +257,7 @@ def test_shape(df_raw: Any) -> None: @pytest.mark.parametrize("df_raw", [df_polars, df_pandas, df_mpd, df_lazy]) -def test_expr(df_raw: Any) -> None: +def test_expr_binary(df_raw: Any) -> None: result = nw.LazyFrame(df_raw).with_columns( a=(1 + 3 * nw.col("a")) * (1 / nw.col("a")), b=nw.col("z") / (2 - nw.col("b")), @@ -281,3 +281,42 @@ def test_expr(df_raw: Any) -> None: "f": [1, 1, 1], } compare_dicts(result_native, expected) + + +@pytest.mark.parametrize("df_raw", [df_polars, df_pandas, df_lazy]) +def test_expr_unary(df_raw: Any) -> None: + result = ( + nw.LazyFrame(df_raw) + .with_columns( + a_mean=nw.col("a").mean(), + a_sum=nw.col("a").sum(), + b_nunique=nw.col("b").n_unique(), + z_min=nw.col("z").min(), + z_max=nw.col("z").max(), + ) + .select(nw.col("a_mean", "a_sum", "b_nunique", "z_min", "z_max").unique()) + ) + result_native = nw.to_native(result) + expected = {"a_mean": [2], "a_sum": [6], "b_nunique": [2], "z_min": [7], "z_max": [9]} + compare_dicts(result_native, expected) + + +@pytest.mark.parametrize("df_raw", [df_polars, df_pandas, df_mpd, df_lazy]) +def test_expr_transform(df_raw: Any) -> None: + result = nw.LazyFrame(df_raw).with_columns( + a=nw.col("a").is_between(-1, 1), b=nw.col("b").is_in([4, 5]) + ) + result_native = nw.to_native(result) + expected = {"a": [True, False, False], "b": [True, True, False], "z": [7, 8, 9]} + compare_dicts(result_native, expected) + + +@pytest.mark.parametrize("df_raw", [df_polars, df_pandas, df_lazy]) +def test_expr_min_max(df_raw: Any) -> None: + df = nw.LazyFrame(df_raw) + result_min = nw.to_native(df.select(nw.min("a", "b", "z"))) + result_max = nw.to_native(df.select(nw.max("a", "b", "z"))) + expected_min = {"a": [1], "b": [4], "z": [7]} + expected_max = {"a": [3], "b": [6], "z": [9]} + compare_dicts(result_min, expected_min) + compare_dicts(result_max, expected_max)