Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: dask expr dunder methods #740

Merged
merged 1 commit into from
Aug 7, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
104 changes: 104 additions & 0 deletions narwhals/_dask/expr.py
Original file line number Diff line number Diff line change
Expand Up @@ -149,6 +149,14 @@ def __add__(self, other: Any) -> Self:
returns_scalar=False,
)

def __radd__(self, other: Any) -> Self:
return self._from_call(
lambda _input, other: _input.__radd__(other),
"__radd__",
other,
returns_scalar=False,
)

def __sub__(self, other: Any) -> Self:
return self._from_call(
lambda _input, other: _input.__sub__(other),
Expand All @@ -157,6 +165,14 @@ def __sub__(self, other: Any) -> Self:
returns_scalar=False,
)

def __rsub__(self, other: Any) -> Self:
return self._from_call(
lambda _input, other: _input.__rsub__(other),
"__rsub__",
other,
returns_scalar=False,
)

def __mul__(self, other: Any) -> Self:
return self._from_call(
lambda _input, other: _input.__mul__(other),
Expand All @@ -165,6 +181,78 @@ def __mul__(self, other: Any) -> Self:
returns_scalar=False,
)

def __rmul__(self, other: Any) -> Self:
return self._from_call(
lambda _input, other: _input.__rmul__(other),
"__rmul__",
other,
returns_scalar=False,
)

def __truediv__(self, other: Any) -> Self:
return self._from_call(
lambda _input, other: _input.__truediv__(other),
"__truediv__",
other,
returns_scalar=False,
)

def __rtruediv__(self, other: Any) -> Self:
return self._from_call(
lambda _input, other: _input.__rtruediv__(other),
"__rtruediv__",
other,
returns_scalar=False,
)

def __floordiv__(self, other: Any) -> Self:
return self._from_call(
lambda _input, other: _input.__floordiv__(other),
"__floordiv__",
other,
returns_scalar=False,
)

def __rfloordiv__(self, other: Any) -> Self:
return self._from_call(
lambda _input, other: _input.__rfloordiv__(other),
"__rfloordiv__",
other,
returns_scalar=False,
)

def __pow__(self, other: Any) -> Self:
return self._from_call(
lambda _input, other: _input.__pow__(other),
"__pow__",
other,
returns_scalar=False,
)

def __rpow__(self, other: Any) -> Self:
return self._from_call(
lambda _input, other: _input.__rpow__(other),
"__rpow__",
other,
returns_scalar=False,
)

def __mod__(self, other: Any) -> Self:
return self._from_call(
lambda _input, other: _input.__mod__(other),
"__mod__",
other,
returns_scalar=False,
)

def __rmod__(self, other: Any) -> Self:
return self._from_call(
lambda _input, other: _input.__rmod__(other),
"__rmod__",
other,
returns_scalar=False,
)

def __eq__(self, other: DaskExpr) -> Self: # type: ignore[override]
return self._from_call(
lambda _input, other: _input.__eq__(other),
Expand Down Expand Up @@ -221,6 +309,14 @@ def __and__(self, other: DaskExpr) -> Self:
returns_scalar=False,
)

def __rand__(self, other: DaskExpr) -> Self: # pragma: no cover
return self._from_call(
lambda _input, other: _input.__rand__(other),
"__rand__",
other,
returns_scalar=False,
)

def __or__(self, other: DaskExpr) -> Self:
return self._from_call(
lambda _input, other: _input.__or__(other),
Expand All @@ -229,6 +325,14 @@ def __or__(self, other: DaskExpr) -> Self:
returns_scalar=False,
)

def __ror__(self, other: DaskExpr) -> Self: # pragma: no cover
return self._from_call(
lambda _input, other: _input.__ror__(other),
"__ror__",
other,
returns_scalar=False,
)

def mean(self) -> Self:
return self._from_call(
lambda _input: _input.mean(),
Expand Down
8 changes: 0 additions & 8 deletions tests/expr_and_series/arithmetic_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,12 +28,6 @@ def test_arithmetic(
constructor: Any,
request: Any,
) -> None:
if "dask" in str(constructor) and attr not in [
"__add__",
"__sub__",
"__mul__",
]:
request.applymarker(pytest.mark.xfail)
if attr == "__mod__" and any(
x in str(constructor) for x in ["pandas_pyarrow", "pyarrow_table", "modin"]
):
Expand Down Expand Up @@ -64,8 +58,6 @@ def test_right_arithmetic(
constructor: Any,
request: Any,
) -> None:
if "dask" in str(constructor):
request.applymarker(pytest.mark.xfail)
if attr == "__rmod__" and any(
x in str(constructor) for x in ["pandas_pyarrow", "pyarrow_table", "modin"]
):
Expand Down
Loading