Repository: dask-contrib/dask-expr
Pull request #898 (2024-02-26 17:35:12)

Revert "Use Aggregation from dask/dask"
Reverts dask-contrib/dask-expr#895

Should revert this so that we can release on top of the current released dask version


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 26cf4ba..23dfddd 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
 ###
@@ -393,18 +453,24 @@ class HolisticGroupbyAggregation(GroupbyAggregationBase):
 
     This class always calculates the aggregates by first collecting all the data for
     the groups and then aggregating at once.
+
+    We are always shuffling, so we will never call combine
     """
 
+    @functools.cached_property
+    def _meta(self):
+        meta = self._meta_chunk
+        aggregate = self.aggregate or (lambda x: x)
+        aggregate_kwargs = self.aggregate_kwargs
+        meta = aggregate([meta], **aggregate_kwargs)
+        return make_meta(meta)
+
     chunk = staticmethod(_non_agg_chunk)
 
     @property
     def should_shuffle(self):
         return True
 
-    @classmethod
-    def combine(cls, inputs, **kwargs):
-        return _groupby_aggregate_spec(_concat(inputs), **kwargs)
-
     @classmethod
     def aggregate(cls, inputs, **kwargs):
         return _groupby_aggregate_spec(_concat(inputs), **kwargs)
@@ -418,15 +484,6 @@ class HolisticGroupbyAggregation(GroupbyAggregationBase):
             **_as_dict("dropna", self.dropna),
         }
 
-    @property
-    def combine_kwargs(self) -> dict:
-        return {
-            "spec": self.arg,
-            "levels": _determine_levels(self.by),
-            **_as_dict("observed", self.observed),
-            **_as_dict("dropna", self.dropna),
-        }
-
     @property
     def aggregate_kwargs(self) -> dict:
         return {


Test patch:
diff --git a/dask_expr/tests/test_groupby.py b/dask_expr/tests/test_groupby.py
index 4d1a088..d5d064d 100644
--- a/dask_expr/tests/test_groupby.py
+++ b/dask_expr/tests/test_groupby.py
@@ -4,10 +4,9 @@ from collections import OrderedDict
 import dask
 import numpy as np
 import pytest
-from dask.dataframe import Aggregation
 
 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
@@ -955,3 +954,9 @@ def test_groupby_size_drop_columns(df, pdf):
     result = df.groupby("x").size()
     assert_eq(result, pdf.groupby("x").size())
     assert result.simplify()._name == df[[]].groupby("x").size().simplify()._name
+
+
+def test_groupby_agg_meta_error(df, pdf):
+    result = df.groupby(["x"]).agg({"y": ["median", "std"]})
+    expected = pdf.groupby(["x"]).agg({"y": ["median", "std"]})
+    assert_eq(result, expected)