Repository: narwhals-dev/narwhals
Pull request #718 (2024-08-04 21:51:37)

feat: add `group_by` to dask
## What type of PR is this? (check all applicable)

- [ ] 💾 Refactor
- [x] ✨ Feature
- [ ] 🐛 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📝 Documentation
- [x] ✅ Test
- [ ] 🐳 Other

## Related issues 

- Related issue #637
- Closes #

## Checklist

- [x] Code follows style guide (ruff)
- [x] Tests added 
- [ ] Documented the changes

## If you have comments or can explain your changes, please do so below.

Mostly copied code from `PandasLikeGroupBy`. Updated tests to `collect()` dask dataframes before sorting.


Patch:
diff --git a/docs/api-reference/expr.md b/docs/api-reference/expr.md
index 492066aa..d37a1266 100644
--- a/docs/api-reference/expr.md
+++ b/docs/api-reference/expr.md
@@ -18,6 +18,7 @@
         - filter
         - gather_every
         - head
+        - clip
         - is_between
         - is_duplicated
         - is_first_distinct
diff --git a/docs/api-reference/series.md b/docs/api-reference/series.md
index 89ed25c6..011e4357 100644
--- a/docs/api-reference/series.md
+++ b/docs/api-reference/series.md
@@ -20,6 +20,7 @@
         - gather_every
         - head
         - is_between
+        - clip
         - is_duplicated
         - is_empty
         - is_first_distinct
diff --git a/narwhals/_dask/dataframe.py b/narwhals/_dask/dataframe.py
index 11cefb05..b7256319 100644
--- a/narwhals/_dask/dataframe.py
+++ b/narwhals/_dask/dataframe.py
@@ -123,3 +123,8 @@ class DaskLazyFrame:
 
     def rename(self: Self, mapping: dict[str, str]) -> Self:
         return self._from_native_dataframe(self._native_dataframe.rename(columns=mapping))
+
+    def group_by(self, *by: str) -> Any:
+        from narwhals._dask.group_by import DaskLazyGroupBy
+
+        return DaskLazyGroupBy(self, list(by))
diff --git a/narwhals/_dask/group_by.py b/narwhals/_dask/group_by.py
new file mode 100644
index 00000000..69f11c12
--- /dev/null
+++ b/narwhals/_dask/group_by.py
@@ -0,0 +1,146 @@
+from __future__ import annotations
+
+import collections
+from copy import copy
+from typing import TYPE_CHECKING
+from typing import Any
+from typing import Callable
+
+from narwhals._expression_parsing import is_simple_aggregation
+from narwhals._expression_parsing import parse_into_exprs
+from narwhals.utils import remove_prefix
+
+if TYPE_CHECKING:
+    from narwhals._dask.dataframe import DaskLazyFrame
+    from narwhals._dask.expr import DaskExpr
+    from narwhals._dask.typing import IntoDaskExpr
+
+POLARS_TO_PANDAS_AGGREGATIONS = {
+    "len": "size",
+}
+
+
+class DaskLazyGroupBy:
+    def __init__(self, df: DaskLazyFrame, keys: list[str]) -> None:
+        self._df = df
+        self._keys = keys
+        self._grouped = self._df._native_dataframe.groupby(
+            list(self._keys),
+            dropna=False,
+        )
+
+    def agg(
+        self,
+        *aggs: IntoDaskExpr,
+        **named_aggs: IntoDaskExpr,
+    ) -> DaskLazyFrame:
+        exprs = parse_into_exprs(
+            *aggs,
+            namespace=self._df.__narwhals_namespace__(),
+            **named_aggs,
+        )
+        output_names: list[str] = copy(self._keys)
+        for expr in exprs:
+            # TODO(unassigned): Need to implement `.all()` for Dask first, part of GH 637
+            if expr._output_names is None:  # pragma: no cover
+                msg = (
+                    "Anonymous expressions are not supported in group_by.agg.\n"
+                    "Instead of `nw.all()`, try using a named expression, such as "
+                    "`nw.col('a', 'b')`\n"
+                )
+                raise ValueError(msg)
+
+            output_names.extend(expr._output_names)
+
+        return agg_dask(
+            self._grouped,
+            exprs,
+            self._keys,
+            output_names,
+            self._from_native_dataframe,
+        )
+
+    def _from_native_dataframe(self, df: DaskLazyFrame) -> DaskLazyFrame:
+        from narwhals._dask.dataframe import DaskLazyFrame
+
+        return DaskLazyFrame(
+            df,
+            backend_version=self._df._backend_version,
+        )
+
+
+def agg_dask(
+    grouped: Any,
+    exprs: list[DaskExpr],
+    keys: list[str],
+    output_names: list[str],
+    from_dataframe: Callable[[Any], DaskLazyFrame],
+) -> DaskLazyFrame:
+    """
+    This should be the fastpath, but cuDF is too far behind to use it.
+
+    - https://github.com/rapidsai/cudf/issues/15118
+    - https://github.com/rapidsai/cudf/issues/15084
+    """
+    all_simple_aggs = True
+    for expr in exprs:
+        if not is_simple_aggregation(expr):
+            all_simple_aggs = False
+            break
+
+    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
+                # 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"
+                    raise AssertionError(msg)
+
+                function_name = POLARS_TO_PANDAS_AGGREGATIONS.get(
+                    expr._function_name, expr._function_name
+                )
+                for output_name in expr._output_names:
+                    simple_aggregations[output_name] = (keys[0], function_name)
+                continue
+
+            # e.g. agg(nw.mean('a')) # noqa: ERA001
+            if (
+                expr._depth != 1 or expr._root_names is None or expr._output_names is None
+            ):  # pragma: no cover
+                msg = "Safety assertion failed, please report a bug to https://github.com/narwhals-dev/narwhals/issues"
+                raise AssertionError(msg)
+
+            function_name = remove_prefix(expr._function_name, "col->")
+            function_name = POLARS_TO_PANDAS_AGGREGATIONS.get(
+                function_name, function_name
+            )
+            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)
+        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])
+
+    msg = (
+        "Non-trivial complex found.\n\n"
+        "Hint: you were probably trying to apply a non-elementary aggregation with a "
+        "dask dataframe.\n"
+        "Please rewrite your query such that group-by aggregations "
+        "are elementary. For example, instead of:\n\n"
+        "    df.group_by('a').agg(nw.col('b').round(2).mean())\n\n"
+        "use:\n\n"
+        "    df.with_columns(nw.col('b').round(2)).group_by('a').agg(nw.col('b').mean())\n\n"
+    )
+    raise ValueError(msg)
diff --git a/narwhals/_pandas_like/expr.py b/narwhals/_pandas_like/expr.py
index e798bd4d..5231412e 100644
--- a/narwhals/_pandas_like/expr.py
+++ b/narwhals/_pandas_like/expr.py
@@ -196,6 +196,12 @@ class PandasLikeExpr:
         return reuse_series_implementation(self, "min", returns_scalar=True)
 
     # Other
+
+    def clip(self, lower_bound: Any, upper_bound: Any) -> Self:
+        return reuse_series_implementation(
+            self, "clip", lower_bound=lower_bound, upper_bound=upper_bound
+        )
+
     def is_between(
         self, lower_bound: Any, upper_bound: Any, closed: str = "both"
     ) -> Self:
diff --git a/narwhals/_pandas_like/series.py b/narwhals/_pandas_like/series.py
index e0746186..acd99bf5 100644
--- a/narwhals/_pandas_like/series.py
+++ b/narwhals/_pandas_like/series.py
@@ -631,6 +631,13 @@ class PandasLikeSeries:
     def gather_every(self: Self, n: int, offset: int = 0) -> Self:
         return self._from_native_series(self._native_series.iloc[offset::n])
 
+    def clip(
+        self: Self, lower_bound: Any | None = None, upper_bound: Any | None = None
+    ) -> Self:
+        return self._from_native_series(
+            self._native_series.clip(lower_bound, upper_bound)
+        )
+
     @property
     def str(self) -> PandasLikeSeriesStringNamespace:
         return PandasLikeSeriesStringNamespace(self)
diff --git a/narwhals/expr.py b/narwhals/expr.py
index baab0f25..27185a6e 100644
--- a/narwhals/expr.py
+++ b/narwhals/expr.py
@@ -1765,6 +1765,122 @@ class Expr:
             lambda plx: self._call(plx).gather_every(n=n, offset=offset)
         )
 
+    # need to allow numeric typing
+    # TODO @aivanoved: make type alias for numeric type
+    def clip(
+        self,
+        lower_bound: IntoExpr | Any | None = None,
+        upper_bound: IntoExpr | Any | None = None,
+    ) -> Self:
+        r"""
+        Clip values in the Series.
+
+        Arguments:
+            lower_bound: Lower bound value.
+            upper_bound: Upper bound value.
+
+        Examples:
+            >>> import pandas as pd
+            >>> import polars as pl
+            >>> import narwhals as nw
+            >>>
+            >>> s = [1, 2, 3]
+            >>> df_pd = pd.DataFrame({"s": s})
+            >>> df_pl = pl.DataFrame({"s": s})
+
+            We define a library agnostic function:
+
+            >>> @nw.narwhalify
+            ... def func_lower(df):
+            ...     return df.select(nw.col("s").clip(2))
+
+            We can then pass either pandas or Polars to `func_lower`:
+
+            >>> func_lower(df_pd)
+               s
+            0  2
+            1  2
+            2  3
+            >>> func_lower(df_pl)  # doctest: +NORMALIZE_WHITESPACE
+            shape: (3, 1)
+            ┌─────┐
+            │ s   │
+            │ --- │
+            │ i64 │
+            ╞═════╡
+            │ 2   │
+            │ 2   │
+            │ 3   │
+            └─────┘
+
+            We define another library agnostic function:
+
+            >>> @nw.narwhalify
+            ... def func_upper(df):
+            ...     return df.select(nw.col("s").clip(upper_bound=2))
+
+            We can then pass either pandas or Polars to `func_upper`:
+
+            >>> func_upper(df_pd)
+               s
+            0  1
+            1  2
+            2  2
+            >>> func_upper(df_pl)  # doctest: +NORMALIZE_WHITESPACE
+            shape: (3, 1)
+            ┌─────┐
+            │ s   │
+            │ --- │
+            │ i64 │
+            ╞═════╡
+            │ 1   │
+            │ 2   │
+            │ 2   │
+            └─────┘
+
+            We can have both at the same time
+
+            >>> s = [-1, 1, -3, 3, -5, 5]
+            >>> df_pd = pd.DataFrame({"s": s})
+            >>> df_pl = pl.DataFrame({"s": s})
+
+            We define a library agnostic function:
+
+            >>> @nw.narwhalify
+            ... def func(df):
+            ...     return df.select(nw.col("s").clip(-1, 3))
+
+            We can pass either pandas or Polars to `func`:
+
+            >>> func(df_pd)
+               s
+            0 -1
+            1  1
+            2 -1
+            3  3
+            4 -1
+            5  3
+            >>> func(df_pl)  # doctest: +NORMALIZE_WHITESPACE
+            shape: (6, 1)
+            ┌─────┐
+            │ s   │
+            │ --- │
+            │ i64 │
+            ╞═════╡
+            │ -1  │
+            │ 1   │
+            │ -1  │
+            │ 3   │
+            │ -1  │
+            │ 3   │
+            └─────┘
+        """
+        return self.__class__(
+            lambda plx: self._call(plx).clip(
+                extract_compliant(plx, lower_bound), extract_compliant(plx, upper_bound)
+            )
+        )
+
     @property
     def str(self: Self) -> ExprStringNamespace:
         return ExprStringNamespace(self)
diff --git a/narwhals/series.py b/narwhals/series.py
index 2582978c..b8e98cd7 100644
--- a/narwhals/series.py
+++ b/narwhals/series.py
@@ -544,6 +544,107 @@ class Series:
         """
         return self._compliant_series.std(ddof=ddof)
 
+    def clip(
+        self, lower_bound: Any | None = None, upper_bound: Any | None = None
+    ) -> Self:
+        r"""
+        Clip values in the Series.
+
+        Arguments:
+            lower_bound: Lower bound value.
+            upper_bound: Upper bound value.
+
+        Examples:
+            >>> import pandas as pd
+            >>> import polars as pl
+            >>> import narwhals as nw
+            >>>
+            >>> s = [1, 2, 3]
+            >>> s_pd = pd.Series(s)
+            >>> s_pl = pl.Series(s)
+
+            We define a library agnostic function:
+
+            >>> @nw.narwhalify
+            ... def func_lower(s):
+            ...     return s.clip(2)
+
+            We can then pass either pandas or Polars to `func_lower`:
+
+            >>> func_lower(s_pd)
+            0    2
+            1    2
+            2    3
+            dtype: int64
+            >>> func_lower(s_pl)  # doctest: +NORMALIZE_WHITESPACE
+            shape: (3,)
+            Series: '' [i64]
+            [
+               2
+               2
+               3
+            ]
+
+            We define another library agnostic function:
+
+            >>> @nw.narwhalify
+            ... def func_upper(s):
+            ...     return s.clip(upper_bound=2)
+
+            We can then pass either pandas or Polars to `func_upper`:
+
+            >>> func_upper(s_pd)
+            0    1
+            1    2
+            2    2
+            dtype: int64
+            >>> func_upper(s_pl)  # doctest: +NORMALIZE_WHITESPACE
+            shape: (3,)
+            Series: '' [i64]
+            [
+               1
+               2
+               2
+            ]
+
+            We can have both at the same time
+
+            >>> s = [-1, 1, -3, 3, -5, 5]
+            >>> s_pd = pd.Series(s)
+            >>> s_pl = pl.Series(s)
+
+            We define a library agnostic function:
+
+            >>> @nw.narwhalify
+            ... def func(s):
+            ...     return s.clip(-1, 3)
+
+            We can pass either pandas or Polars to `func`:
+
+            >>> func(s_pd)
+            0   -1
+            1    1
+            2   -1
+            3    3
+            4   -1
+            5    3
+            dtype: int64
+            >>> func(s_pl)  # doctest: +NORMALIZE_WHITESPACE
+            shape: (6,)
+            Series: '' [i64]
+            [
+               -1
+               1
+               -1
+               3
+               -1
+               3
+            ]
+        """
+        return self._from_compliant_series(
+            self._compliant_series.clip(lower_bound=lower_bound, upper_bound=upper_bound)
+        )
+
     def is_in(self, other: Any) -> Self:
         """
         Check if the elements of this Series are in the other sequence.
diff --git a/pyproject.toml b/pyproject.toml
index b9f72b67..63f54661 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -91,6 +91,7 @@ filterwarnings = [
   'ignore:.*implementation has mismatches with pandas',
   'ignore:.*Do not use the `random` module inside strategies',
   'ignore:.*You are using pyarrow version',
+  'ignore:.*but when imported by pytest it raised',
 ]
 xfail_strict = true
 markers = ["slow: marks tests as slow (deselect with '-m \"not slow\"')"]


Test patch:
diff --git a/tests/expr_and_series/clip_test.py b/tests/expr_and_series/clip_test.py
new file mode 100644
index 00000000..94b54272
--- /dev/null
+++ b/tests/expr_and_series/clip_test.py
@@ -0,0 +1,24 @@
+from typing import Any
+
+import pytest
+
+import narwhals.stable.v1 as nw
+from tests.utils import compare_dicts
+
+
+def test_clip(request: Any, constructor: Any) -> None:
+    if "pyarrow_table" in str(constructor) or "dask" in str(constructor):
+        request.applymarker(pytest.mark.xfail)
+    df = nw.from_native(constructor({"a": [1, 2, 3, -4, 5]}))
+    result = df.select(b=nw.col("a").clip(3, 5))
+    expected = {"b": [3, 3, 3, 3, 5]}
+    compare_dicts(result, expected)
+
+
+def test_clip_series(request: Any, constructor_eager: Any) -> None:
+    if "pyarrow_table" in str(constructor_eager):
+        request.applymarker(pytest.mark.xfail)
+    df = nw.from_native(constructor_eager({"a": [1, 2, 3, -4, 5]}), eager_only=True)
+    result = {"b": df["a"].clip(3, 5)}
+    expected = {"b": [3, 3, 3, 3, 5]}
+    compare_dicts(result, expected)
diff --git a/tests/test_group_by.py b/tests/test_group_by.py
index 20537185..53d595c3 100644
--- a/tests/test_group_by.py
+++ b/tests/test_group_by.py
@@ -35,6 +35,20 @@ def test_group_by_complex() -> None:
     compare_dicts(result, expected)
 
 
+def test_invalid_group_by_dask() -> None:
+    pytest.importorskip("dask")
+    pytest.importorskip("dask_expr")
+    import dask.dataframe as dd
+
+    df_dask = dd.from_pandas(df_pandas)
+
+    with pytest.raises(ValueError, match=r"Non-trivial complex found"):
+        nw.from_native(df_dask).group_by("a").agg(nw.col("b").mean().min())
+
+    with pytest.raises(RuntimeError, match="does your"):
+        nw.from_native(df_dask).group_by("a").agg(nw.col("b"))
+
+
 def test_invalid_group_by() -> None:
     df = nw.from_native(df_pandas)
     with pytest.raises(RuntimeError, match="does your"):
@@ -66,11 +80,11 @@ def test_group_by_iter(constructor_eager: Any) -> None:
     assert sorted(keys) == sorted(expected_keys)
     expected_keys = [(1, 4), (3, 6)]  # type: ignore[list-item]
     keys = []
-    for key, _df in df.group_by("a", "b"):
+    for key, _ in df.group_by("a", "b"):
         keys.append(key)
     assert sorted(keys) == sorted(expected_keys)
     keys = []
-    for key, _df in df.group_by(["a", "b"]):
+    for key, _ in df.group_by(["a", "b"]):
         keys.append(key)
     assert sorted(keys) == sorted(expected_keys)
 
@@ -94,17 +108,16 @@ def test_group_by_empty_result_pandas() -> None:
         )
 
 
-def test_group_by_simple_named(constructor: Any, request: Any) -> None:
-    if "dask" in str(constructor):
-        request.applymarker(pytest.mark.xfail)
+def test_group_by_simple_named(constructor: Any) -> None:
     data = {"a": [1, 1, 2], "b": [4, 5, 6], "c": [7, 2, 1]}
-    df = nw.from_native(constructor(data))
+    df = nw.from_native(constructor(data)).lazy()
     result = (
         df.group_by("a")
         .agg(
             b_min=nw.col("b").min(),
             b_max=nw.col("b").max(),
         )
+        .collect()
         .sort("a")
     )
     expected = {
@@ -115,17 +128,16 @@ def test_group_by_simple_named(constructor: Any, request: Any) -> None:
     compare_dicts(result, expected)
 
 
-def test_group_by_simple_unnamed(constructor: Any, request: Any) -> None:
-    if "dask" in str(constructor):
-        request.applymarker(pytest.mark.xfail)
+def test_group_by_simple_unnamed(constructor: Any) -> None:
     data = {"a": [1, 1, 2], "b": [4, 5, 6], "c": [7, 2, 1]}
-    df = nw.from_native(constructor(data))
+    df = nw.from_native(constructor(data)).lazy()
     result = (
         df.group_by("a")
         .agg(
             nw.col("b").min(),
             nw.col("c").max(),
         )
+        .collect()
         .sort("a")
     )
     expected = {
@@ -136,17 +148,16 @@ def test_group_by_simple_unnamed(constructor: Any, request: Any) -> None:
     compare_dicts(result, expected)
 
 
-def test_group_by_multiple_keys(constructor: Any, request: Any) -> None:
-    if "dask" in str(constructor):
-        request.applymarker(pytest.mark.xfail)
+def test_group_by_multiple_keys(constructor: Any) -> None:
     data = {"a": [1, 1, 2], "b": [4, 4, 6], "c": [7, 2, 1]}
-    df = nw.from_native(constructor(data))
+    df = nw.from_native(constructor(data)).lazy()
     result = (
         df.group_by("a", "b")
         .agg(
             c_min=nw.col("c").min(),
             c_max=nw.col("c").max(),
         )
+        .collect()
         .sort("a")
     )
     expected = {