From f6bcf459d4be2d5bcb1139db89a632205340f647 Mon Sep 17 00:00:00 2001 From: FBruzzesi Date: Sun, 11 Aug 2024 11:05:39 +0200 Subject: [PATCH] feat: dask sum_horizontal --- narwhals/_arrow/namespace.py | 18 +++--------------- narwhals/_dask/namespace.py | 3 +++ narwhals/_pandas_like/namespace.py | 18 +++--------------- tests/expr_and_series/sum_horizontal_test.py | 4 +--- 4 files changed, 10 insertions(+), 33 deletions(-) diff --git a/narwhals/_arrow/namespace.py b/narwhals/_arrow/namespace.py index fbb285b50..57bd5a4f1 100644 --- a/narwhals/_arrow/namespace.py +++ b/narwhals/_arrow/namespace.py @@ -167,25 +167,13 @@ def _lit_arrow_series(_: ArrowDataFrame) -> ArrowSeries: ) def all_horizontal(self, *exprs: IntoArrowExpr) -> ArrowExpr: - return reduce( - lambda x, y: x & y, - parse_into_exprs(*exprs, namespace=self), - ) + return reduce(lambda x, y: x & y, parse_into_exprs(*exprs, namespace=self)) def any_horizontal(self, *exprs: IntoArrowExpr) -> ArrowExpr: - return reduce( - lambda x, y: x | y, - parse_into_exprs(*exprs, namespace=self), - ) + return reduce(lambda x, y: x | y, parse_into_exprs(*exprs, namespace=self)) def sum_horizontal(self, *exprs: IntoArrowExpr) -> ArrowExpr: - return reduce( - lambda x, y: x + y, - parse_into_exprs( - *exprs, - namespace=self, - ), - ) + return reduce(lambda x, y: x + y, parse_into_exprs(*exprs, namespace=self)) def concat( self, diff --git a/narwhals/_dask/namespace.py b/narwhals/_dask/namespace.py index 82ade973b..82a0e1586 100644 --- a/narwhals/_dask/namespace.py +++ b/narwhals/_dask/namespace.py @@ -133,6 +133,9 @@ def all_horizontal(self, *exprs: IntoDaskExpr) -> DaskExpr: def any_horizontal(self, *exprs: IntoDaskExpr) -> DaskExpr: return reduce(lambda x, y: x | y, parse_into_exprs(*exprs, namespace=self)) + def sum_horizontal(self, *exprs: IntoDaskExpr) -> DaskExpr: + return reduce(lambda x, y: x + y, parse_into_exprs(*exprs, namespace=self)) + def _create_expr_from_series(self, _: Any) -> NoReturn: msg = "`_create_expr_from_series` for DaskNamespace exists only for compatibility" raise NotImplementedError(msg) diff --git a/narwhals/_pandas_like/namespace.py b/narwhals/_pandas_like/namespace.py index 154e333be..13e2e99d3 100644 --- a/narwhals/_pandas_like/namespace.py +++ b/narwhals/_pandas_like/namespace.py @@ -202,25 +202,13 @@ def len(self) -> PandasLikeExpr: # --- horizontal --- def sum_horizontal(self, *exprs: IntoPandasLikeExpr) -> PandasLikeExpr: - return reduce( - lambda x, y: x + y, - parse_into_exprs( - *exprs, - namespace=self, - ), - ) + return reduce(lambda x, y: x + y, parse_into_exprs(*exprs, namespace=self)) def all_horizontal(self, *exprs: IntoPandasLikeExpr) -> PandasLikeExpr: - return reduce( - lambda x, y: x & y, - parse_into_exprs(*exprs, namespace=self), - ) + return reduce(lambda x, y: x & y, parse_into_exprs(*exprs, namespace=self)) def any_horizontal(self, *exprs: IntoPandasLikeExpr) -> PandasLikeExpr: - return reduce( - lambda x, y: x | y, - parse_into_exprs(*exprs, namespace=self), - ) + return reduce(lambda x, y: x | y, parse_into_exprs(*exprs, namespace=self)) def concat( self, diff --git a/tests/expr_and_series/sum_horizontal_test.py b/tests/expr_and_series/sum_horizontal_test.py index dd9c5d906..9411903cb 100644 --- a/tests/expr_and_series/sum_horizontal_test.py +++ b/tests/expr_and_series/sum_horizontal_test.py @@ -7,9 +7,7 @@ @pytest.mark.parametrize("col_expr", [nw.col("a"), "a"]) -def test_sumh(constructor: Any, col_expr: Any, request: Any) -> None: - if "dask" in str(constructor): - request.applymarker(pytest.mark.xfail) +def test_sumh(constructor: Any, col_expr: Any) -> None: data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]} df = nw.from_native(constructor(data)) result = df.with_columns(horizontal_sum=nw.sum_horizontal(col_expr, nw.col("b")))