Repository: dask/dask

Issue:
Implement groupby/median
In #9483 @jrbourbeau added `median` and `median_approximate` to dask.dataframe. It would also be nice to have `DataFrame.groupby(...).median()`. This is challenging to do with an ACA-based approach for the same reason that `DataFrame.median()` is. However, it could be tractable with a shuffle-based approach (cf. #9487)



Discussion:

Resolving patch:
diff --git a/dask/dataframe/groupby.py b/dask/dataframe/groupby.py
--- a/dask/dataframe/groupby.py
+++ b/dask/dataframe/groupby.py
@@ -341,6 +341,76 @@ def _groupby_aggregate(
     return aggfunc(grouped, **kwargs)
 
 
+def _groupby_aggregate_spec(
+    df, spec, levels=None, dropna=None, sort=False, observed=None, **kwargs
+):
+    """
+    A simpler version of _groupby_aggregate that just calls ``aggregate`` using
+    the user-provided spec.
+    """
+    dropna = {"dropna": dropna} if dropna is not None else {}
+    observed = {"observed": observed} if observed is not None else {}
+    return df.groupby(level=levels, sort=sort, **observed, **dropna).aggregate(
+        spec, **kwargs
+    )
+
+
+def _non_agg_chunk(df, *by, key, dropna=None, observed=None, **kwargs):
+    """
+    A non-aggregation agg function. This simuates the behavior of an initial
+    partitionwise aggregation, but doesn't actually aggregate or throw away
+    any data.
+    """
+    if is_series_like(df):
+        # Handle a series-like groupby. `by` could be columns that are not the series,
+        # but are like-indexed, so we handle that case by temporarily converting to
+        # a dataframe, then setting the index.
+        result = df.to_frame().set_index(by[0] if len(by) == 1 else list(by))[df.name]
+    else:
+        # Handle a frame-like groupby.
+        result = df.set_index(list(by))
+        if isinstance(key, (tuple, list, set, pd.Index)):
+            key = list(key)
+        result = result[key]
+
+    # If observed is False, we have to check for categorical indices and possibly enrich
+    # them with unobserved values. This function is intended as an initial partitionwise
+    # aggregation, so you might expect that we could enrich the frame with unobserved
+    # values at the end. However, if you have multiple output partitions, that results
+    # in duplicated unobserved values in each partition. So we have to do this step
+    # at the start before any shuffling occurs so that we can consolidate all of the
+    # unobserved values in a single partition.
+    if observed is False:
+        has_categoricals = False
+
+        # Search for categorical indices and get new index objects that have all the
+        # categories in them.
+        if isinstance(result.index, pd.CategoricalIndex):
+            has_categoricals = True
+            full_index = result.index.categories.copy().rename(result.index.name)
+        elif isinstance(result.index, pd.MultiIndex) and any(
+            isinstance(level, pd.CategoricalIndex) for level in result.index.levels
+        ):
+            has_categoricals = True
+            full_index = pd.MultiIndex.from_product(
+                (
+                    level.categories
+                    if isinstance(level, pd.CategoricalIndex)
+                    else level
+                    for level in result.index.levels
+                ),
+                names=result.index.names,
+            )
+        if has_categoricals:
+            # If we found any categoricals, append unobserved values to the end of the
+            # frame.
+            new_cats = full_index[~full_index.isin(result.index)]
+            empty = pd.DataFrame(pd.NA, index=new_cats, columns=result.columns)
+            result = pd.concat([result, empty])
+
+    return result
+
+
 def _apply_chunk(df, *by, dropna=None, observed=None, **kwargs):
     func = kwargs.pop("chunk")
     columns = kwargs.pop("columns")
@@ -733,7 +803,7 @@ def _build_agg_args(spec):
         applied after the ``agg_funcs``. They are used to create final results
         from intermediate representations.
     """
-    known_np_funcs = {np.min: "min", np.max: "max"}
+    known_np_funcs = {np.min: "min", np.max: "max", np.median: "median"}
 
     # check that there are no name conflicts for a single input column
     by_name = {}
@@ -779,6 +849,10 @@ def _build_agg_args_single(result_column, func, input_column):
         "first": (M.first, M.first),
         "last": (M.last, M.last),
         "prod": (M.prod, M.prod),
+        "median": (
+            None,
+            M.median,
+        ),  # No chunk func for median, we can only take it when aggregating
     }
 
     if func in simple_impl.keys():
@@ -1629,6 +1703,41 @@ def mean(self, split_every=None, split_out=1, shuffle=None):
             c = c[s.columns]
         return s / c
 
+    @derived_from(pd.core.groupby.GroupBy)
+    def median(self, split_every=None, split_out=1, shuffle=None):
+        if shuffle is False:
+            raise ValueError(
+                "In order to aggregate with 'median', you must use shuffling-based "
+                "aggregation (e.g., shuffle='tasks')"
+            )
+        shuffle = shuffle or config.get("shuffle", None) or "tasks"
+
+        with check_numeric_only_deprecation():
+            meta = self._meta_nonempty.median()
+        columns = meta.name if is_series_like(meta) else meta.columns
+        by = self.by if isinstance(self.by, list) else [self.by]
+        return _shuffle_aggregate(
+            [self.obj] + by,
+            token="non-agg",
+            chunk=_non_agg_chunk,
+            chunk_kwargs={
+                "key": columns,
+                **self.observed,
+                **self.dropna,
+            },
+            aggregate=_groupby_aggregate,
+            aggregate_kwargs={
+                "aggfunc": _median_aggregate,
+                "levels": _determine_levels(self.by),
+                **self.observed,
+                **self.dropna,
+            },
+            split_every=split_every,
+            split_out=split_out,
+            shuffle=shuffle,
+            sort=self.sort,
+        )
+
     @derived_from(pd.core.groupby.GroupBy)
     def size(self, split_every=None, split_out=1, shuffle=None):
         return self._single_agg(
@@ -1860,6 +1969,15 @@ def aggregate(self, arg, split_every=None, split_out=1, shuffle=None):
                 f"if pandas < 1.1.0. Pandas version is {pd.__version__}"
             )
 
+        # If any of the agg funcs contain a "median", we *must* use the shuffle
+        # implementation.
+        has_median = any(s[1] in ("median", np.median) for s in spec)
+        if has_median and not shuffle:
+            raise ValueError(
+                "In order to aggregate with 'median', you must use shuffling-based "
+                "aggregation (e.g., shuffle='tasks')"
+            )
+
         if shuffle:
             # Shuffle-based aggregation
             #
@@ -1867,29 +1985,56 @@ def aggregate(self, arg, split_every=None, split_out=1, shuffle=None):
             # for larger values of split_out. However, the shuffle
             # step requires that the result of `chunk` produces a
             # proper DataFrame type
-            return _shuffle_aggregate(
-                chunk_args,
-                chunk=_groupby_apply_funcs,
-                chunk_kwargs=dict(
-                    funcs=chunk_funcs,
-                    sort=False,
-                    **self.observed,
-                    **self.dropna,
-                ),
-                aggregate=_agg_finalize,
-                aggregate_kwargs=dict(
-                    aggregate_funcs=aggregate_funcs,
-                    finalize_funcs=finalizers,
-                    level=levels,
-                    **self.observed,
-                    **self.dropna,
-                ),
-                token="aggregate",
-                split_every=split_every,
-                split_out=split_out,
-                shuffle=shuffle if isinstance(shuffle, str) else "tasks",
-                sort=self.sort,
-            )
+            # If we have a median in the spec, we cannot do an initial
+            # aggregation.
+            if has_median:
+                return _shuffle_aggregate(
+                    chunk_args,
+                    chunk=_non_agg_chunk,
+                    chunk_kwargs={
+                        "key": [
+                            c for c in _obj.columns.tolist() if c not in group_columns
+                        ],
+                        **self.observed,
+                        **self.dropna,
+                    },
+                    aggregate=_groupby_aggregate_spec,
+                    aggregate_kwargs={
+                        "spec": arg,
+                        "levels": _determine_levels(self.by),
+                        **self.observed,
+                        **self.dropna,
+                    },
+                    token="aggregate",
+                    split_every=split_every,
+                    split_out=split_out,
+                    shuffle=shuffle if isinstance(shuffle, str) else "tasks",
+                    sort=self.sort,
+                )
+            else:
+                return _shuffle_aggregate(
+                    chunk_args,
+                    chunk=_groupby_apply_funcs,
+                    chunk_kwargs={
+                        "funcs": chunk_funcs,
+                        "sort": self.sort,
+                        **self.observed,
+                        **self.dropna,
+                    },
+                    aggregate=_agg_finalize,
+                    aggregate_kwargs=dict(
+                        aggregate_funcs=aggregate_funcs,
+                        finalize_funcs=finalizers,
+                        level=levels,
+                        **self.observed,
+                        **self.dropna,
+                    ),
+                    token="aggregate",
+                    split_every=split_every,
+                    split_out=split_out,
+                    shuffle=shuffle if isinstance(shuffle, str) else "tasks",
+                    sort=self.sort,
+                )
 
         if self.sort is None and split_out > 1:
             warnings.warn(SORT_SPLIT_OUT_WARNING, FutureWarning)
@@ -2552,6 +2697,11 @@ def _head_aggregate(series_gb, **kwargs):
     return series_gb.head(**kwargs).droplevel(list(range(levels)))
 
 
+def _median_aggregate(series_gb, **kwargs):
+    with check_numeric_only_deprecation():
+        return series_gb.median(**kwargs)
+
+
 def _shuffle_aggregate(
     args,
     chunk=None,
@@ -2625,13 +2775,15 @@ def _shuffle_aggregate(
     elif split_every < 1 or not isinstance(split_every, Integral):
         raise ValueError("split_every must be an integer >= 1")
 
-    # Shuffle-based groupby aggregation
+    # Shuffle-based groupby aggregation. N.B. we have to use `_meta_nonempty`
+    # as some chunk aggregation functions depend on the index having values to
+    # determine `group_keys` behavior.
     chunk_name = f"{token or funcname(chunk)}-chunk"
     chunked = map_partitions(
         chunk,
         *args,
         meta=chunk(
-            *[arg._meta if isinstance(arg, _Frame) else arg for arg in args],
+            *[arg._meta_nonempty if isinstance(arg, _Frame) else arg for arg in args],
             **chunk_kwargs,
         ),
         token=chunk_name,


Test patch:
diff --git a/dask/dataframe/tests/test_groupby.py b/dask/dataframe/tests/test_groupby.py
--- a/dask/dataframe/tests/test_groupby.py
+++ b/dask/dataframe/tests/test_groupby.py
@@ -30,6 +30,7 @@
 AGG_FUNCS = [
     "sum",
     "mean",
+    "median",
     "min",
     "max",
     "count",
@@ -1000,7 +1001,8 @@ def test_aggregate__single_element_groups(agg_func):
     if spec in {"mean", "var"}:
         expected = expected.astype(float)
 
-    assert_eq(expected, ddf.groupby(["a", "d"]).agg(spec))
+    shuffle = {"shuffle": "tasks"} if agg_func == "median" else {}
+    assert_eq(expected, ddf.groupby(["a", "d"]).agg(spec, **shuffle))
 
 
 def test_aggregate_build_agg_args__reuse_of_intermediates():
@@ -1207,6 +1209,29 @@ def test_shuffle_aggregate_defaults(shuffle_method):
     assert any("shuffle" in l for l in dsk.layers)
 
 
+@pytest.mark.parametrize("spec", [{"c": "median"}, {"b": np.median, "c": np.max}])
+@pytest.mark.parametrize("keys", ["a", ["a", "d"]])
+def test_aggregate_median(spec, keys, shuffle_method):
+    pdf = pd.DataFrame(
+        {
+            "a": [1, 2, 3, 1, 1, 2, 4, 3, 7] * 10,
+            "b": [4, 2, 7, 3, 3, 1, 1, 1, 2] * 10,
+            "c": [0, 1, 2, 3, 4, 5, 6, 7, 8] * 10,
+            "d": [3, 2, 1, 3, 2, 1, 2, 6, 4] * 10,
+        },
+        columns=["c", "b", "a", "d"],
+    )
+    ddf = dd.from_pandas(pdf, npartitions=10)
+    actual = ddf.groupby(keys).aggregate(spec, shuffle=shuffle_method)
+    expected = pdf.groupby(keys).aggregate(spec)
+    assert_eq(actual, expected)
+
+    with pytest.raises(ValueError, match="must use shuffl"):
+        ddf.groupby(keys).aggregate(spec, shuffle=False)
+    with pytest.raises(ValueError, match="must use shuffl"):
+        ddf.groupby(keys).median(shuffle=False)
+
+
 @pytest.mark.parametrize("axis", [0, 1])
 @pytest.mark.parametrize("group_keys", [True, False, None])
 @pytest.mark.parametrize("method", ["ffill", "bfill"])