Skip to content

Commit

Permalink
feat: dask sum_horizontal (#775)
Browse files Browse the repository at this point in the history
* feat: dask sum_horizontal
  • Loading branch information
FBruzzesi authored Aug 11, 2024
1 parent 35e33d1 commit 75d61e4
Show file tree
Hide file tree
Showing 5 changed files with 12 additions and 35 deletions.
18 changes: 3 additions & 15 deletions narwhals/_arrow/namespace.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
3 changes: 3 additions & 0 deletions narwhals/_dask/namespace.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
18 changes: 3 additions & 15 deletions narwhals/_pandas_like/namespace.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
4 changes: 1 addition & 3 deletions tests/expr_and_series/sum_horizontal_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")))
Expand Down
4 changes: 2 additions & 2 deletions tests/test_selectors.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
import pytest

import narwhals.stable.v1 as nw
from narwhals.dependencies import get_dask
from narwhals.dependencies import get_dask_dataframe
from narwhals.selectors import all
from narwhals.selectors import boolean
from narwhals.selectors import by_dtype
Expand Down Expand Up @@ -67,7 +67,7 @@ def test_categorical(request: Any, constructor: Any) -> None:
compare_dicts(result, expected)


@pytest.mark.skipif((get_dask() is None), reason="too old for dask")
@pytest.mark.skipif((get_dask_dataframe() is None), reason="too old for dask")
def test_dask_categorical() -> None:
import dask.dataframe as dd

Expand Down

0 comments on commit 75d61e4

Please sign in to comment.