refactor: Implement DaskLazyGroupBy

Much happier with this than the `pandas` one

Author: dangotbanned

narwhals/_dask/dataframe.py

@@ -388,7 +388,7 @@ def join_asof(
     def group_by(self: Self, *by: str, drop_null_keys: bool) -> DaskLazyGroupBy:
         from narwhals._dask.group_by import DaskLazyGroupBy
 
-        return DaskLazyGroupBy(self, list(by), drop_null_keys=drop_null_keys)
+        return DaskLazyGroupBy(self, by, drop_null_keys=drop_null_keys)
 
     def tail(self: Self, n: int) -> Self:  # pragma: no cover
         native_frame = self._native_frame

narwhals/_dask/group_by.py

@@ -5,13 +5,14 @@
 from typing import TYPE_CHECKING
 from typing import Any
 from typing import Callable
+from typing import ClassVar
 from typing import Mapping
 from typing import Sequence
 
 import dask.dataframe as dd
 
+from narwhals._compliant import CompliantGroupBy
 from narwhals._expression_parsing import evaluate_output_names_and_aliases
-from narwhals._expression_parsing import is_elementary_expression
 
 try:
     import dask.dataframe.dask_expr as dx
@@ -54,90 +55,54 @@ def std(ddof: int) -> _AggFn:
     return partial(_DaskGroupBy.std, ddof=ddof)
 
 
-POLARS_TO_DASK_AGGREGATIONS: Mapping[str, Aggregation] = {
-    "sum": "sum",
-    "mean": "mean",
-    "median": "median",
-    "max": "max",
-    "min": "min",
-    "std": std,
-    "var": var,
-    "len": "size",
-    "n_unique": n_unique,
-    "count": "count",
-}
+class DaskLazyGroupBy(CompliantGroupBy["DaskLazyFrame", "DaskExpr"]):
+    _NARWHALS_TO_NATIVE_AGGREGATIONS: ClassVar[Mapping[str, Aggregation]] = {
+        "sum": "sum",
+        "mean": "mean",
+        "median": "median",
+        "max": "max",
+        "min": "min",
+        "std": std,
+        "var": var,
+        "len": "size",
+        "n_unique": n_unique,
+        "count": "count",
+    }
 
-
-class DaskLazyGroupBy:
     def __init__(
-        self: Self, df: DaskLazyFrame, keys: list[str], *, drop_null_keys: bool
+        self: Self, df: DaskLazyFrame, keys: Sequence[str], /, *, drop_null_keys: bool
     ) -> None:
-        self._df: DaskLazyFrame = df
-        self._keys = keys
-        self._grouped = self._df._native_frame.groupby(
-            list(self._keys),
-            dropna=drop_null_keys,
-            observed=True,
-        )
-
-    def agg(
-        self: Self,
-        *exprs: DaskExpr,
-    ) -> DaskLazyFrame:
-        return agg_dask(
-            self._df,
-            self._grouped,
-            exprs,
-            self._keys,
-            self._from_native_frame,
+        self._compliant_frame = df
+        self._keys: list[str] = list(keys)
+        self._grouped = self.compliant.native.groupby(
+            list(self._keys), dropna=drop_null_keys, observed=True
         )
 
-    def _from_native_frame(self: Self, df: dd.DataFrame) -> DaskLazyFrame:
+    def agg(self: Self, *exprs: DaskExpr) -> DaskLazyFrame:
         from narwhals._dask.dataframe import DaskLazyFrame
 
-        return DaskLazyFrame(
-            df,
-            backend_version=self._df._backend_version,
-            version=self._df._version,
-        )
-
-
-def agg_dask(
-    df: DaskLazyFrame,
-    grouped: Any,
-    exprs: Sequence[DaskExpr],
-    keys: 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
-    """
-    if not exprs:
-        # No aggregation provided
-        return df.simple_select(*keys).unique(subset=keys, keep="any")
-
-    all_simple_aggs = True
-    for expr in exprs:
-        if not (
-            is_elementary_expression(expr)
-            and re.sub(r"(\w+->)", "", expr._function_name) in POLARS_TO_DASK_AGGREGATIONS
-        ):
-            all_simple_aggs = False
-            break
-
-    if all_simple_aggs:
+        if not exprs:
+            # No aggregation provided
+            return self.compliant.simple_select(*self._keys).unique(
+                self._keys, keep="any"
+            )
+        self._ensure_all_simple(exprs)
+        # 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
+        POLARS_TO_DASK_AGGREGATIONS = self._NARWHALS_TO_NATIVE_AGGREGATIONS  # noqa: N806
         simple_aggregations: dict[str, tuple[str, Aggregation]] = {}
         for expr in exprs:
-            output_names, aliases = evaluate_output_names_and_aliases(expr, df, keys)
+            output_names, aliases = evaluate_output_names_and_aliases(
+                expr, self.compliant, self._keys
+            )
             if expr._depth == 0:
                 # e.g. agg(nw.len()) # noqa: ERA001
                 function_name = POLARS_TO_DASK_AGGREGATIONS.get(
                     expr._function_name, expr._function_name
                 )
                 simple_aggregations.update(
-                    dict.fromkeys(aliases, (keys[0], function_name))
+                    dict.fromkeys(aliases, (self._keys[0], function_name))
                 )
                 continue
 
@@ -150,24 +115,12 @@ def agg_dask(
                 if callable(agg_function)
                 else agg_function
             )
-
             simple_aggregations.update(
-                {
-                    alias: (output_name, agg_function)
-                    for alias, output_name in zip(aliases, output_names)
-                }
+                (alias, (output_name, agg_function))
+                for alias, output_name in zip(aliases, output_names)
             )
-        result_simple = grouped.agg(**simple_aggregations)
-        return from_dataframe(result_simple.reset_index())
-
-    msg = (
-        "Non-trivial complex aggregation 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)
+        return DaskLazyFrame(
+            self._grouped.agg(**simple_aggregations).reset_index(),
+            backend_version=self.compliant._backend_version,
+            version=self.compliant._version,
+        )