Repository: narwhals-dev/narwhals
Pull request #618 (2024-07-24 21:56:04)

feat: arrow `__mod__` method
## What type of PR is this? (check all applicable)

- [ ] 💾 Refactor
- [x] ✨ Feature
- [ ] 🐛 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📝 Documentation
- [ ] ✅ Test
- [ ] 🐳 Other

## Related issues 

- Related issue #372 

## Checklist

- [x] Code follows style guide (ruff)
- [x] Tests added 
- [x] Documented the changes

## If you have comments or can explain your changes, please do so below.

Added test for the concerns in the [comment of #491](https://github.com/narwhals-dev/narwhals/pull/491#issuecomment-2225707517)

FBruzzesi: Was there some change in pandas floordiv?!
FBruzzesi: Strong statement, but after a quick look, it seems that this is the last pure pyarrow_table xfail in the tests. Meaning that the remaining ones are either for version specific or behavior difference 👀

If so, should we close #372?

Patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 0091cc0c..757cff1d 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -31,7 +31,7 @@ repos:
       additional_dependencies: [polars]
     - id: imports-are-banned
       name: import are banned (use `get_pandas` instead of `import pandas`)
-      entry: (?<!>>> )import (pandas|polars|modin|cudf|pyarrow)
+      entry: (?<!>>> )import (pandas|polars|modin|cudf|pyarrow|dask)
       language: pygrep
       files: ^narwhals/
       exclude: ^narwhals/dependencies\.py
diff --git a/narwhals/_arrow/expr.py b/narwhals/_arrow/expr.py
index 25400dc7..ac0fbbcc 100644
--- a/narwhals/_arrow/expr.py
+++ b/narwhals/_arrow/expr.py
@@ -145,6 +145,12 @@ class ArrowExpr:
     def __rtruediv__(self, other: ArrowExpr | Any) -> Self:
         return reuse_series_implementation(self, "__rtruediv__", other)
 
+    def __mod__(self, other: ArrowExpr | Any) -> Self:
+        return reuse_series_implementation(self, "__mod__", other)
+
+    def __rmod__(self, other: ArrowExpr | Any) -> Self:
+        return reuse_series_implementation(self, "__rmod__", other)
+
     def __invert__(self) -> Self:
         return reuse_series_implementation(self, "__invert__")
 
diff --git a/narwhals/_arrow/series.py b/narwhals/_arrow/series.py
index bd50c030..fadc72b8 100644
--- a/narwhals/_arrow/series.py
+++ b/narwhals/_arrow/series.py
@@ -190,6 +190,22 @@ class ArrowSeries:
             other = pa.scalar(other)
         return self._from_native_series(pc.divide(*cast_for_truediv(other, ser)))
 
+    def __mod__(self, other: Any) -> Self:
+        pc = get_pyarrow_compute()
+        ser = self._native_series
+        other = validate_column_comparand(other)
+        floor_div = (self // other)._native_series
+        res = pc.subtract(ser, pc.multiply(floor_div, other))
+        return self._from_native_series(res)
+
+    def __rmod__(self, other: Any) -> Self:
+        pc = get_pyarrow_compute()
+        ser = self._native_series
+        other = validate_column_comparand(other)
+        floor_div = (other // self)._native_series
+        res = pc.subtract(other, pc.multiply(floor_div, ser))
+        return self._from_native_series(res)
+
     def __invert__(self) -> Self:
         pc = get_pyarrow_compute()
         return self._from_native_series(pc.invert(self._native_series))
diff --git a/narwhals/_dask/expr.py b/narwhals/_dask/expr.py
index 27a574a6..91b2c461 100644
--- a/narwhals/_dask/expr.py
+++ b/narwhals/_dask/expr.py
@@ -4,6 +4,7 @@ from copy import copy
 from typing import TYPE_CHECKING
 from typing import Any
 from typing import Callable
+from typing import NoReturn
 
 from narwhals.dependencies import get_dask
 
@@ -405,6 +406,13 @@ class DaskExpr:
             returns_scalar=True,
         )
 
+    def count(self) -> Self:
+        return self._from_call(
+            lambda _input: _input.count(),
+            "count",
+            returns_scalar=True,
+        )
+
     def round(self, decimals: int) -> Self:
         return self._from_call(
             lambda _input, decimals: _input.round(decimals),
@@ -413,6 +421,23 @@ class DaskExpr:
             returns_scalar=False,
         )
 
+    def drop_nulls(self) -> NoReturn:
+        # We can't (yet?) allow methods which modify the index
+        msg = "`Expr.drop_nulls` is not supported for the Dask backend. Please use `LazyFrame.drop_nulls` instead."
+        raise NotImplementedError(msg)
+
+    def sort(self, *, descending: bool = False, nulls_last: bool = False) -> NoReturn:
+        # We can't (yet?) allow methods which modify the index
+        msg = "`Expr.sort` is not supported for the Dask backend. Please use `LazyFrame.sort` instead."
+        raise NotImplementedError(msg)
+
+    def abs(self) -> Self:
+        return self._from_call(
+            lambda _input: _input.abs(),
+            "abs",
+            returns_scalar=False,
+        )
+
     def fill_null(self, value: Any) -> DaskExpr:
         return self._from_call(
             lambda _input, _val: _input.fillna(_val),
diff --git a/narwhals/_dask/group_by.py b/narwhals/_dask/group_by.py
index 32a43002..57e112c8 100644
--- a/narwhals/_dask/group_by.py
+++ b/narwhals/_dask/group_by.py
@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import collections
 from copy import copy
 from typing import TYPE_CHECKING
 from typing import Any
@@ -55,7 +54,6 @@ class DaskLazyGroupBy:
             self._grouped,
             exprs,
             self._keys,
-            output_names,
             self._from_native_dataframe,
         )
 
@@ -72,7 +70,6 @@ def agg_dask(
     grouped: Any,
     exprs: list[DaskExpr],
     keys: list[str],
-    output_names: list[str],
     from_dataframe: Callable[[Any], DaskLazyFrame],
 ) -> DaskLazyFrame:
     """
@@ -90,8 +87,7 @@ def agg_dask(
     if all_simple_aggs:
         simple_aggregations: dict[str, tuple[str, str]] = {}
         for expr in exprs:
-            # TODO(unassigned): Need to implement `.len()` for Dask first
-            if expr._depth == 0:  # pragma: no cover
+            if expr._depth == 0:
                 # e.g. agg(nw.len()) # noqa: ERA001
                 if expr._output_names is None:  # pragma: no cover
                     msg = "Safety assertion failed, please report a bug to https://github.com/narwhals-dev/narwhals/issues"
@@ -117,20 +113,12 @@ def agg_dask(
             )
             for root_name, output_name in zip(expr._root_names, expr._output_names):
                 simple_aggregations[output_name] = (root_name, function_name)
-
-        aggs = collections.defaultdict(list)
-        name_mapping = {}
-        for output_name, named_agg in simple_aggregations.items():
-            aggs[named_agg[0]].append(named_agg[1])
-            name_mapping[f"{named_agg[0]}_{named_agg[1]}"] = output_name
         try:
-            result_simple = grouped.agg(aggs)
+            result_simple = grouped.agg(**simple_aggregations)
         except ValueError as exc:
             msg = "Failed to aggregated - does your aggregation function return a scalar?"
             raise RuntimeError(msg) from exc
-        result_simple.columns = [f"{a}_{b}" for a, b in result_simple.columns]
-        result_simple = result_simple.rename(columns=name_mapping).reset_index()
-        return from_dataframe(result_simple.loc[:, output_names])
+        return from_dataframe(result_simple.reset_index())
 
     msg = (
         "Non-trivial complex found.\n\n"
diff --git a/narwhals/_dask/namespace.py b/narwhals/_dask/namespace.py
index fdf7092c..600f9d82 100644
--- a/narwhals/_dask/namespace.py
+++ b/narwhals/_dask/namespace.py
@@ -9,6 +9,8 @@ from typing import NoReturn
 from narwhals import dtypes
 from narwhals._dask.expr import DaskExpr
 from narwhals._expression_parsing import parse_into_exprs
+from narwhals.dependencies import get_dask_dataframe
+from narwhals.dependencies import get_pandas
 
 if TYPE_CHECKING:
     from narwhals._dask.dataframe import DaskLazyFrame
@@ -85,6 +87,23 @@ class DaskNamespace:
             backend_version=self._backend_version,
         ).sum()
 
+    def len(self) -> DaskExpr:
+        dd = get_dask_dataframe()
+        pd = get_pandas()
+
+        # coverage bug? this is definitely hit
+        return DaskExpr(  # pragma: no cover
+            lambda df: [
+                dd.from_pandas(pd.Series([len(df._native_dataframe)], name="len"))["len"]
+            ],
+            depth=0,
+            function_name="len",
+            root_names=None,
+            output_names=["len"],
+            returns_scalar=True,
+            backend_version=self._backend_version,
+        )
+
     def all_horizontal(self, *exprs: IntoDaskExpr) -> DaskExpr:
         return reduce(lambda x, y: x & y, parse_into_exprs(*exprs, namespace=self))
 


Test patch:
diff --git a/tests/expr_and_series/abs_test.py b/tests/expr_and_series/abs_test.py
index c612b31f..e684528b 100644
--- a/tests/expr_and_series/abs_test.py
+++ b/tests/expr_and_series/abs_test.py
@@ -1,14 +1,10 @@
 from typing import Any
 
-import pytest
-
 import narwhals.stable.v1 as nw
 from tests.utils import compare_dicts
 
 
-def test_abs(constructor: Any, request: Any) -> None:
-    if "dask" in str(constructor):
-        request.applymarker(pytest.mark.xfail)
+def test_abs(constructor: Any) -> None:
     df = nw.from_native(constructor({"a": [1, 2, 3, -4, 5]}))
     result = df.select(b=nw.col("a").abs())
     expected = {"b": [1, 2, 3, 4, 5]}
diff --git a/tests/expr_and_series/arithmetic_test.py b/tests/expr_and_series/arithmetic_test.py
index fbddb6f1..0f17b1f2 100644
--- a/tests/expr_and_series/arithmetic_test.py
+++ b/tests/expr_and_series/arithmetic_test.py
@@ -29,7 +29,7 @@ def test_arithmetic(
     request: Any,
 ) -> None:
     if attr == "__mod__" and any(
-        x in str(constructor) for x in ["pandas_pyarrow", "pyarrow_table", "modin"]
+        x in str(constructor) for x in ["pandas_pyarrow", "modin"]
     ):
         request.applymarker(pytest.mark.xfail)
 
@@ -59,7 +59,7 @@ def test_right_arithmetic(
     request: Any,
 ) -> None:
     if attr == "__rmod__" and any(
-        x in str(constructor) for x in ["pandas_pyarrow", "pyarrow_table", "modin"]
+        x in str(constructor) for x in ["pandas_pyarrow", "modin"]
     ):
         request.applymarker(pytest.mark.xfail)
 
@@ -90,7 +90,7 @@ def test_arithmetic_series(
     request: Any,
 ) -> None:
     if attr == "__mod__" and any(
-        x in str(constructor_eager) for x in ["pandas_pyarrow", "pyarrow_table", "modin"]
+        x in str(constructor_eager) for x in ["pandas_pyarrow", "modin"]
     ):
         request.applymarker(pytest.mark.xfail)
 
@@ -120,7 +120,7 @@ def test_right_arithmetic_series(
     request: Any,
 ) -> None:
     if attr == "__rmod__" and any(
-        x in str(constructor_eager) for x in ["pandas_pyarrow", "pyarrow_table", "modin"]
+        x in str(constructor_eager) for x in ["pandas_pyarrow", "modin"]
     ):
         request.applymarker(pytest.mark.xfail)
 
diff --git a/tests/expr_and_series/count_test.py b/tests/expr_and_series/count_test.py
index 10feff5d..208df3bc 100644
--- a/tests/expr_and_series/count_test.py
+++ b/tests/expr_and_series/count_test.py
@@ -1,14 +1,10 @@
 from typing import Any
 
-import pytest
-
 import narwhals.stable.v1 as nw
 from tests.utils import compare_dicts
 
 
-def test_count(constructor: Any, request: Any) -> None:
-    if "dask" in str(constructor):
-        request.applymarker(pytest.mark.xfail)
+def test_count(constructor: Any) -> None:
     data = {"a": [1, 3, 2], "b": [4, None, 6], "z": [7.0, None, None]}
     df = nw.from_native(constructor(data))
     result = df.select(nw.col("a", "b", "z").count())
diff --git a/tests/series_only/arithmetic_test.py b/tests/series_only/arithmetic_test.py
index b9eee678..9219ad4f 100644
--- a/tests/series_only/arithmetic_test.py
+++ b/tests/series_only/arithmetic_test.py
@@ -35,7 +35,7 @@ def test_truediv_same_dims(constructor_eager: Any, request: Any) -> None:
 @pytest.mark.skipif(
     parse_version(pd.__version__) < (2, 0), reason="convert_dtypes not available"
 )
-def test_mod(left: int, right: int) -> None:
+def test_floordiv(left: int, right: int) -> None:
     # hypothesis complains if we add `constructor` as an argument, so this
     # test is a bit manual unfortunately
     assume(right != 0)
@@ -65,3 +65,34 @@ def test_mod(left: int, right: int) -> None:
         nw.col("a") // right
     )
     compare_dicts(result, expected)
+
+
+@pytest.mark.slow()
+@given(  # type: ignore[misc]
+    left=st.integers(-100, 100),
+    right=st.integers(-100, 100),
+)
+@pytest.mark.skipif(
+    parse_version(pd.__version__) < (2, 0), reason="convert_dtypes not available"
+)
+def test_mod(left: int, right: int) -> None:
+    # hypothesis complains if we add `constructor` as an argument, so this
+    # test is a bit manual unfortunately
+    assume(right != 0)
+    expected = {"a": [left % right]}
+    result = nw.from_native(pd.DataFrame({"a": [left]}), eager_only=True).select(
+        nw.col("a") % right
+    )
+    compare_dicts(result, expected)
+    result = nw.from_native(
+        pd.DataFrame({"a": [left]}).convert_dtypes(), eager_only=True
+    ).select(nw.col("a") % right)
+    compare_dicts(result, expected)
+    result = nw.from_native(pl.DataFrame({"a": [left]}), eager_only=True).select(
+        nw.col("a") % right
+    )
+    compare_dicts(result, expected)
+    result = nw.from_native(pa.table({"a": [left]}), eager_only=True).select(
+        nw.col("a") % right
+    )
+    compare_dicts(result, expected)
diff --git a/tests/tpch_q1_test.py b/tests/tpch_q1_test.py
index da540e46..999f32f7 100644
--- a/tests/tpch_q1_test.py
+++ b/tests/tpch_q1_test.py
@@ -17,7 +17,7 @@ from tests.utils import compare_dicts
 
 @pytest.mark.parametrize(
     "library",
-    ["pandas", "polars", "pyarrow"],
+    ["pandas", "polars", "pyarrow", "dask"],
 )
 @pytest.mark.filterwarnings("ignore:.*Passing a BlockManager.*:DeprecationWarning")
 def test_q1(library: str, request: Any) -> None:
@@ -28,6 +28,14 @@ def test_q1(library: str, request: Any) -> None:
         df_raw["l_shipdate"] = pd.to_datetime(df_raw["l_shipdate"])
     elif library == "polars":
         df_raw = pl.scan_parquet("tests/data/lineitem.parquet")
+    elif library == "dask":
+        pytest.importorskip("dask")
+        pytest.importorskip("dask_expr", exc_type=ImportError)
+        import dask.dataframe as dd
+
+        df_raw = dd.from_pandas(
+            pd.read_parquet("tests/data/lineitem.parquet", dtype_backend="pyarrow")
+        )
     else:
         df_raw = pq.read_table("tests/data/lineitem.parquet")
     var_1 = datetime(1998, 9, 2)