Repository: dask-contrib/dask-expr
Pull request #893 (2024-02-23 11:13:10)

Add support for dd.Aggregation
sibling pr https://github.com/dask/dask/pull/10947


Patch:
diff --git a/dask_expr/__init__.py b/dask_expr/__init__.py
index 0eebc0d..bd602c7 100644
--- a/dask_expr/__init__.py
+++ b/dask_expr/__init__.py
@@ -2,6 +2,7 @@ from dask_expr import _version, datasets
 from dask_expr._collection import *
 from dask_expr._dispatch import get_collection_type
 from dask_expr._dummies import get_dummies
+from dask_expr._groupby import Aggregation
 from dask_expr.io._delayed import from_delayed
 from dask_expr.io.bag import to_bag
 from dask_expr.io.csv import to_csv
diff --git a/dask_expr/_groupby.py b/dask_expr/_groupby.py
index d19787d..11c6181 100644
--- a/dask_expr/_groupby.py
+++ b/dask_expr/_groupby.py
@@ -81,6 +81,66 @@ def _adjust_split_out_for_group_keys(npartitions, by):
     return math.ceil(npartitions / (100 / (len(by) - 1)))
 
 
+class Aggregation:
+    """User defined groupby-aggregation.
+
+    This class allows users to define their own custom aggregation in terms of
+    operations on Pandas dataframes in a map-reduce style. You need to specify
+    what operation to do on each chunk of data, how to combine those chunks of
+    data together, and then how to finalize the result.
+
+    See :ref:`dataframe.groupby.aggregate` for more.
+
+    Parameters
+    ----------
+    name : str
+        the name of the aggregation. It should be unique, since intermediate
+        result will be identified by this name.
+    chunk : callable
+        a function that will be called with the grouped column of each
+        partition. It can either return a single series or a tuple of series.
+        The index has to be equal to the groups.
+    agg : callable
+        a function that will be called to aggregate the results of each chunk.
+        Again the argument(s) will be grouped series. If ``chunk`` returned a
+        tuple, ``agg`` will be called with all of them as individual positional
+        arguments.
+    finalize : callable
+        an optional finalizer that will be called with the results from the
+        aggregation.
+
+    Examples
+    --------
+    We could implement ``sum`` as follows:
+
+    >>> custom_sum = dd.Aggregation(
+    ...     name='custom_sum',
+    ...     chunk=lambda s: s.sum(),
+    ...     agg=lambda s0: s0.sum()
+    ... )  # doctest: +SKIP
+    >>> df.groupby('g').agg(custom_sum)  # doctest: +SKIP
+
+    We can implement ``mean`` as follows:
+
+    >>> custom_mean = dd.Aggregation(
+    ...     name='custom_mean',
+    ...     chunk=lambda s: (s.count(), s.sum()),
+    ...     agg=lambda count, sum: (count.sum(), sum.sum()),
+    ...     finalize=lambda count, sum: sum / count,
+    ... )  # doctest: +SKIP
+    >>> df.groupby('g').agg(custom_mean)  # doctest: +SKIP
+
+    Though of course, both of these are built-in and so you don't need to
+    implement them yourself.
+    """
+
+    def __init__(self, name, chunk, agg, finalize=None):
+        self.chunk = chunk
+        self.agg = agg
+        self.finalize = finalize
+        self.__name__ = name
+
+
 ###
 ### Groupby-aggregation expressions
 ###
@@ -297,19 +357,6 @@ class GroupbyAggregationBase(GroupByApplyConcatApply, GroupByBase):
         "_slice": None,
     }
 
-    @functools.cached_property
-    def _meta(self):
-        meta = meta_nonempty(self.frame._meta)
-        meta = meta.groupby(
-            self._by_meta,
-            **_as_dict("observed", self.observed),
-            **_as_dict("dropna", self.dropna),
-        )
-        if self._slice is not None:
-            meta = meta[self._slice]
-        meta = meta.aggregate(self.arg)
-        return make_meta(meta)
-
     @functools.cached_property
     def spec(self):
         # Converts the `arg` operand into specific
@@ -373,6 +420,10 @@ class GroupbyAggregation(GroupbyAggregationBase):
     or holistic aggregations.
     """
 
+    @functools.cached_property
+    def _meta(self):
+        return self._lower()._meta
+
     @functools.cached_property
     def _is_decomposable(self):
         return not any(s[1] in ("median", np.median) for s in self.spec)
@@ -410,10 +461,6 @@ class HolisticGroupbyAggregation(GroupbyAggregationBase):
     def should_shuffle(self):
         return True
 
-    @classmethod
-    def chunk(cls, df, *by, **kwargs):
-        return _non_agg_chunk(df, *by, **kwargs)
-
     @classmethod
     def combine(cls, inputs, **kwargs):
         return _groupby_aggregate_spec(_concat(inputs), **kwargs)
@@ -492,7 +539,7 @@ class DecomposableGroupbyAggregation(GroupbyAggregationBase):
             "arg": self.arg,
             "columns": self._slice,
             "finalize_funcs": self.agg_args["finalizers"],
-            "is_series": self._meta.ndim == 1,
+            "is_series": self.frame._meta.ndim == 1,
             "level": self.levels,
             "sort": self.sort,
             **_as_dict("observed", self.observed),


Test patch:
diff --git a/dask_expr/tests/test_groupby.py b/dask_expr/tests/test_groupby.py
index e5e08cf..5569d73 100644
--- a/dask_expr/tests/test_groupby.py
+++ b/dask_expr/tests/test_groupby.py
@@ -6,7 +6,7 @@ import numpy as np
 import pytest
 
 from dask_expr import from_pandas
-from dask_expr._groupby import GroupByUDFBlockwise
+from dask_expr._groupby import Aggregation, GroupByUDFBlockwise
 from dask_expr._reductions import TreeReduce
 from dask_expr._shuffle import Shuffle, divisions_lru
 from dask_expr.io import FromPandas
@@ -923,6 +923,16 @@ def test_series_groupby_not_supported(df, attr):
         getattr(g.x, attr)()
 
 
+def test_dataframe_groupby_agg_custom_sum():
+    pandas_spec = {"b": "sum"}
+    dask_spec = {"b": Aggregation("sum", lambda s: s.sum(), lambda s0: s0.sum())}
+    df = pd.DataFrame({"g": [0, 0, 1] * 3, "b": [1, 2, 3] * 3})
+    ddf = from_pandas(df, npartitions=2)
+    expected = df.groupby("g").aggregate(pandas_spec)
+    result = ddf.groupby("g").aggregate(dask_spec)
+    assert_eq(result, expected)
+
+
 def test_groupby_size_drop_columns(df, pdf):
     result = df.groupby("x").size()
     assert result.simplify()._name == df[["x"]].groupby("x").size().simplify()._name