Repository: dask/dask

Issue:
User defined aggregations
Dask.dataframe has efficient grouped aggregations for known reductions like sum, max, min.
To do this is uses the apply-concat-apply pattern.  We could expose this pattern out to users to help support their own custom aggregations.  This likely just means renaming the currently private `_aca_agg` method, adding a docstring, and adding some documentation.



Discussion:
The signature of the function is:

> def _aca_agg(self, token, func, aggfunc=None, **kwargs)

Can we derive the `token` from the name of the given `func`? None of the rest of public `dask.dataframe` API methods expose the token attribute. We can use `dask.base.tokenize` for this.

So my proposal is to create a new public `aca_agg` function that in turn calls `_aca_agg`.

I believe this is wrong https://github.com/dask/dask/blob/master/dask/dataframe/core.py#L2073 (edited):

```
def apply_concat_apply(args, chunk=None, aggregate=None,
                       columns=no_default, token=None, chunk_kwargs=None,
                       aggregate_kwargs=None, **kwargs):
    """ Apply a function to blocks, the concat, then apply again"""
    if chunk_kwargs is None:
        chunk_kwargs = dict()
    if aggregate_kwargs is None:
        aggregate_kwargs = dict()
    chunk_kwargs.update(kwargs)
    aggregate_kwargs.update(kwargs)

    if not isinstance(args, (tuple, list)):
        args = [args]

    ...

    a = '{0}--first-{1}'.format(token, token_key)
    if len(args) == 1 and isinstance(args[0], _Frame) and not chunk_kwargs:
        dsk = dict(((a, i), (chunk, key))
                   for i, key in enumerate(args[0]._keys()))
    else:
        dsk = dict(((a, i), (apply, chunk, [(x._name, i)
                                            if isinstance(x, _Frame)
                                            else x for x in args],
                                    kwargs))
```

Last `kwargs` should be `chunk_kwargs`. Otherwise the argument `chunk_kwargs` is ignored and only the extra `kwargs` are used.

> So my proposal is to create a new public aca_agg function that in turn calls _aca_agg

This generally sounds good to me.  I recommend rethinking the name `aca_agg` if we intend to expose this to users.  

I recommend looking at using both `dask.utils.funcname` and `dask.base.tokenize` when determining the name.

Your analysis of `chunk_kwargs` seems accurate to me.


Resolving patch:
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -343,7 +343,7 @@ def drop_duplicates(self, **kwargs):
                    token='drop-duplicates')
 
     def __len__(self):
-        return reduction(self, len, np.sum, token='len', meta=int).compute()
+        return self.reduction(len, np.sum, token='len', meta=int).compute()
 
     @insert_meta_param_description(pad=12)
     def map_partitions(self, func, *args, **kwargs):
@@ -415,6 +415,106 @@ def map_partitions(self, func, *args, **kwargs):
         """
         return map_partitions(func, self, *args, **kwargs)
 
+    @insert_meta_param_description(pad=12)
+    def reduction(self, chunk, aggregate=None, meta=no_default,
+                  token=None, chunk_kwargs=None, aggregate_kwargs=None,
+                  **kwargs):
+        """Generic row-wise reductions.
+
+        Parameters
+        ----------
+        chunk : callable
+            Function to operate on each partition. Should return a
+            ``pandas.DataFrame``, ``pandas.Series``, or a scalar.
+        aggregate : callable, optional
+            Function to operate on the concatenated result of ``chunk``. If not
+            specified, defaults to ``chunk``.
+
+            The input to ``aggregate`` depends on the output of ``chunk``.
+            If the output of ``chunk`` is a:
+            - scalar: Input is a Series, with one row per partition.
+            - Series: Input is a DataFrame, with one row per partition. Columns
+              are the rows in the output series.
+            - DataFrame: Input is a DataFrame, with one row per partition.
+              Columns are the columns in the output dataframes.
+
+            Should return a ``pandas.DataFrame``, ``pandas.Series``, or a
+            scalar.
+        $META
+        token : str, optional
+            The name to use for the output keys.
+        chunk_kwargs : dict, optional
+            Keyword arguments to pass on to ``chunk`` only.
+        aggregate_kwargs : dict, optional
+            Keyword arguments to pass on to ``aggregate`` only.
+        kwargs :
+            All remaining keywords will be passed to both ``chunk`` and
+            ``aggregate``.
+
+        Examples
+        --------
+        >>> import pandas as pd
+        >>> import dask.dataframe as dd
+        >>> df = pd.DataFrame({'x': range(50), 'y': range(50, 100)})
+        >>> ddf = dd.from_pandas(df, npartitions=4)
+
+        Count the number of rows in a DataFrame. To do this, count the number
+        of rows in each partition, then sum the results:
+
+        >>> res = ddf.reduction(lambda x: x.count(),
+        ...                     aggregate=lambda x: x.sum())
+        >>> res.compute()
+        x    50
+        y    50
+        dtype: int64
+
+        Count the number of rows in a Series with elements greater than or
+        equal to a value (provided via a keyword).
+
+        >>> def count_greater(x, value=0):
+        ...     return (x >= value).sum()
+        >>> res = ddf.x.reduction(count_greater, aggregate=lambda x: x.sum(),
+        ...                       chunk_kwargs={'value': 25})
+        >>> res.compute()
+        25
+
+        Aggregate both the sum and count of a Series at the same time:
+
+        >>> def sum_and_count(x):
+        ...     return pd.Series({'sum': x.sum(), 'count': x.count()})
+        >>> res = ddf.x.reduction(sum_and_count, aggregate=lambda x: x.sum())
+        >>> res.compute()
+        count      50
+        sum      1225
+        dtype: int64
+
+        Doing the same, but for a DataFrame. Here ``chunk`` returns a
+        DataFrame, meaning the input to ``aggregate`` is a DataFrame with an
+        index with non-unique entries for both 'x' and 'y'. We groupby the
+        index, and sum each group to get the final result.
+
+        >>> def sum_and_count(x):
+        ...     return pd.DataFrame({'sum': x.sum(), 'count': x.count()})
+        >>> res = ddf.reduction(sum_and_count,
+        ...                     aggregate=lambda x: x.groupby(level=0).sum())
+        >>> res.compute()
+           count   sum
+        x     50  1225
+        y     50  3725
+        """
+        if aggregate is None:
+            aggregate = chunk
+
+        chunk_kwargs = chunk_kwargs.copy() if chunk_kwargs else {}
+        chunk_kwargs['aca_chunk'] = chunk
+
+        aggregate_kwargs = aggregate_kwargs.copy() if aggregate_kwargs else {}
+        aggregate_kwargs['aca_aggregate'] = aggregate
+
+        return aca(self, chunk=_reduction_chunk, aggregate=_reduction_aggregate,
+                   meta=meta, token=token, chunk_kwargs=chunk_kwargs,
+                   aggregate_kwargs=aggregate_kwargs, **kwargs)
+
     def random_split(self, p, random_state=None):
         """ Pseudorandomly split dataframe into different pieces row-wise
 
@@ -814,10 +914,6 @@ def _get_binary_operator(cls, op, inv=False):
         else:
             return lambda self, other: elemwise(op, self, other)
 
-    def _aca_agg(self, token, func, aggfunc=None, **kwargs):
-        """ Wrapper for aggregations """
-        raise NotImplementedError
-
     def rolling(self, window, min_periods=None, freq=None, center=False,
                 win_type=None, axis=0):
         """Provides rolling transformations.
@@ -866,40 +962,37 @@ def rolling(self, window, min_periods=None, freq=None, center=False,
     def sum(self, axis=None, skipna=True):
         axis = self._validate_axis(axis)
         meta = self._meta_nonempty.sum(axis=axis, skipna=skipna)
+        token = self._token_prefix + 'sum'
         if axis == 1:
-            return map_partitions(_sum, self, meta=meta,
-                                  token=self._token_prefix + 'sum',
-                                  axis=axis, skipna=skipna)
+            return self.map_partitions(_sum, meta=meta, token=token,
+                                       skipna=skipna, axis=axis)
         else:
-            return self._aca_agg(token='sum', func=_sum,
-                                 skipna=skipna, axis=axis,
-                                 meta=meta)
+            return self.reduction(_sum, meta=meta, token=token,
+                                  skipna=skipna, axis=axis)
 
     @derived_from(pd.DataFrame)
     def max(self, axis=None, skipna=True):
         axis = self._validate_axis(axis)
         meta = self._meta_nonempty.max(axis=axis, skipna=skipna)
+        token = self._token_prefix + 'max'
         if axis == 1:
-            return map_partitions(_max, self, meta=meta,
-                                  token=self._token_prefix + 'max',
-                                  skipna=skipna, axis=axis)
+            return self.map_partitions(_max, meta=meta, token=token,
+                                       skipna=skipna, axis=axis)
         else:
-            return self._aca_agg(token='max', func=_max,
-                                 skipna=skipna, axis=axis,
-                                 meta=meta)
+            return self.reduction(_max, meta=meta, token=token,
+                                  skipna=skipna, axis=axis)
 
     @derived_from(pd.DataFrame)
     def min(self, axis=None, skipna=True):
         axis = self._validate_axis(axis)
         meta = self._meta_nonempty.min(axis=axis, skipna=skipna)
+        token = self._token_prefix + 'min'
         if axis == 1:
-            return map_partitions(_min, self, meta=meta,
-                                  token=self._token_prefix + 'min',
-                                  skipna=skipna, axis=axis)
+            return self.map_partitions(_min, meta=meta, token=token,
+                                       skipna=skipna, axis=axis)
         else:
-            return self._aca_agg(token='min', func=_min,
-                                 skipna=skipna, axis=axis,
-                                 meta=meta)
+            return self.reduction(_min, meta=meta, token=token,
+                                  skipna=skipna, axis=axis)
 
     @derived_from(pd.DataFrame)
     def idxmax(self, axis=None, skipna=True):
@@ -932,16 +1025,15 @@ def idxmin(self, axis=None, skipna=True):
     @derived_from(pd.DataFrame)
     def count(self, axis=None):
         axis = self._validate_axis(axis)
+        token = self._token_prefix + 'count'
         if axis == 1:
             meta = self._meta_nonempty.count(axis=axis)
-            return map_partitions(_count, self, meta=meta,
-                                  token=self._token_prefix + 'count',
-                                  axis=axis)
+            return self.map_partitions(_count, meta=meta, token=token,
+                                       axis=axis)
         else:
             meta = self._meta_nonempty.count()
-            return self._aca_agg(token='count', func=_count,
-                                 aggfunc=lambda x: x.sum(),
-                                 meta=meta)
+            return self.reduction(_count, meta=meta, token=token,
+                                  aggregate=_sum)
 
     @derived_from(pd.DataFrame)
     def mean(self, axis=None, skipna=True):
@@ -1284,8 +1376,8 @@ def column_info(self):
 
     @property
     def nbytes(self):
-        return reduction(self, lambda s: s.nbytes, np.sum, token='nbytes',
-                         meta=int)
+        return self.reduction(lambda s: s.nbytes, np.sum, token='nbytes',
+                              meta=int)
 
     def __array__(self, dtype=None, **kwargs):
         x = np.array(self.compute())
@@ -1351,16 +1443,6 @@ def _validate_axis(cls, axis=0):
         # convert to numeric axis
         return {None: 0, 'index': 0}.get(axis, axis)
 
-    def _aca_agg(self, token, func, aggfunc=None, meta=no_default, **kwargs):
-        """ Wrapper for aggregations """
-        if aggfunc is None:
-            aggfunc = func
-
-        return aca([self], chunk=func,
-                   aggregate=lambda x, **kwargs: aggfunc(pd.Series(x), **kwargs),
-                   meta=meta, token=self._token_prefix + token,
-                   **kwargs)
-
     @derived_from(pd.Series)
     def groupby(self, index, **kwargs):
         from dask.dataframe.groupby import SeriesGroupBy
@@ -1664,18 +1746,17 @@ def nunique(self):
 
     @derived_from(pd.Index)
     def max(self):
-        # it doesn't support axis and skipna kwds
-        return self._aca_agg(token='max', func=_max,
-                             meta=self._meta_nonempty.max())
+        return self.reduction(_max, meta=self._meta_nonempty.max(),
+                              token=self._token_prefix + 'max')
 
     @derived_from(pd.Index)
     def min(self):
-        return self._aca_agg(token='min', func=_min,
-                             meta=self._meta_nonempty.min())
+        return self.reduction(_min, meta=self._meta_nonempty.min(),
+                              token=self._token_prefix + 'min')
 
     def count(self):
-        f = lambda x: pd.notnull(x).sum()
-        return reduction(self, f, np.sum, token='index-count', meta=int)
+        return self.reduction(lambda x: pd.notnull(x).sum(),
+                              np.sum, token='index-count', meta=int)
 
 
 class DataFrame(_Frame):
@@ -2015,20 +2096,6 @@ def _validate_axis(cls, axis=0):
         # convert to numeric axis
         return {None: 0, 'index': 0, 'columns': 1}.get(axis, axis)
 
-    def _aca_agg(self, token, func, aggfunc=None, meta=no_default, **kwargs):
-        """ Wrapper for aggregations """
-        if aggfunc is None:
-            aggfunc = func
-
-        def aggregate(x, **kwargs):
-            return x.groupby(level=0).apply(aggfunc, **kwargs)
-
-        # groupby.aggregation doesn't support skipna,
-        # using gropuby.apply(aggfunc) is a workaround to handle each group as df
-        return aca([self], chunk=func, aggregate=aggregate,
-                   meta=meta, token=self._token_prefix + token,
-                   **kwargs)
-
     @derived_from(pd.DataFrame)
     def drop(self, labels, axis=0, dtype=None):
         if axis != 1:
@@ -2393,33 +2460,6 @@ def empty_safe(func, arg):
         return func(arg)
 
 
-def reduction(x, chunk, aggregate, token=None, meta=no_default):
-    """ General version of reductions
-
-    >>> reduction(my_frame, np.sum, np.sum)  # doctest: +SKIP
-    """
-    token_key = tokenize(x, token or (chunk, aggregate))
-    token = token or 'reduction'
-    a = '{0}--chunk-{1}'.format(token, token_key)
-    dsk = dict(((a, i), (empty_safe, chunk, (x._name, i)))
-               for i in range(x.npartitions))
-
-    b = '{0}--aggregation-{1}'.format(token, token_key)
-    dsk2 = {(b, 0): (aggregate, (remove_empties,
-                     [(a, i) for i in range(x.npartitions)]))}
-
-    if meta is no_default:
-        try:
-            meta_chunk = _emulate(empty_safe, chunk, x)
-            meta = _emulate(aggregate, remove_empties([meta_chunk]))
-        except Exception:
-            raise ValueError("Metadata inference failed, please provide "
-                             "`meta` keyword")
-    meta = make_meta(meta)
-
-    return Scalar(merge(x.dask, dsk, dsk2), b, meta)
-
-
 def _maybe_from_pandas(dfs):
     from .io import from_pandas
     dfs = [from_pandas(df, 1) if isinstance(df, (pd.Series, pd.DataFrame))
@@ -2428,10 +2468,10 @@ def _maybe_from_pandas(dfs):
 
 
 @insert_meta_param_description
-def apply_concat_apply(args, chunk=None, aggregate=None,
-                       meta=no_default, token=None, chunk_kwargs=None,
-                       aggregate_kwargs=None, **kwargs):
-    """ Apply a function to blocks, the concat, then apply again
+def apply_concat_apply(args, chunk=None, aggregate=None, meta=no_default,
+                       token=None, chunk_kwargs=None, aggregate_kwargs=None,
+                       **kwargs):
+    """Apply a function to blocks, then concat, then apply again
 
     Parameters
     ----------
@@ -2476,10 +2516,10 @@ def apply_concat_apply(args, chunk=None, aggregate=None,
     assert all(arg.npartitions == args[0].npartitions
                for arg in args if isinstance(arg, _Frame))
 
-    token_key = tokenize(token or (chunk, aggregate), meta, *args)
-    token = token or 'apply-concat-apply'
+    token_key = tokenize(token or (chunk, aggregate), meta, args,
+                         chunk_kwargs, aggregate_kwargs)
 
-    a = '{0}--first-{1}'.format(token, token_key)
+    a = '{0}-chunk-{1}'.format(token or funcname(chunk), token_key)
     if len(args) == 1 and isinstance(args[0], _Frame) and not chunk_kwargs:
         dsk = dict(((a, i), (chunk, key))
                    for i, key in enumerate(args[0]._keys()))
@@ -2490,7 +2530,7 @@ def apply_concat_apply(args, chunk=None, aggregate=None,
                              chunk_kwargs))
                for i in range(args[0].npartitions))
 
-    b = '{0}--second-{1}'.format(token, token_key)
+    b = '{0}-{1}'.format(token or funcname(aggregate), token_key)
     conc = (_concat, (list, [(a, i) for i in range(args[0].npartitions)]))
     if not aggregate_kwargs:
         dsk2 = {(b, 0): (aggregate, conc)}
@@ -2500,7 +2540,8 @@ def apply_concat_apply(args, chunk=None, aggregate=None,
     if meta is no_default:
         try:
             meta_chunk = _emulate(apply, chunk, args, chunk_kwargs)
-            meta = _emulate(apply, aggregate, [meta_chunk], aggregate_kwargs)
+            meta = _emulate(apply, aggregate, [_concat([meta_chunk])],
+                            aggregate_kwargs)
         except Exception:
             raise ValueError("Metadata inference failed, please provide "
                              "`meta` keyword")
@@ -3229,6 +3270,18 @@ def _std(x, **kwargs):
     return x.std(**kwargs)
 
 
+def _reduction_chunk(x, aca_chunk=None, **kwargs):
+    o = aca_chunk(x, **kwargs)
+    # Return a dataframe so that the concatenated version is also a dataframe
+    return o.to_frame().T if isinstance(o, pd.Series) else o
+
+
+def _reduction_aggregate(x, aca_aggregate=None, **kwargs):
+    if isinstance(x, list):
+        x = pd.Series(x)
+    return aca_aggregate(x, **kwargs)
+
+
 def drop_columns(df, columns, dtype):
     df = df.drop(columns, axis=1)
     df.columns = df.columns.astype(dtype)


Test patch:
diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py
--- a/dask/dataframe/tests/test_dataframe.py
+++ b/dask/dataframe/tests/test_dataframe.py
@@ -12,8 +12,8 @@
 from dask.utils import raises, ignoring, put_lines
 import dask.dataframe as dd
 
-from dask.dataframe.core import (repartition_divisions, _loc, aca, reduction,
-                                 _concat, _Frame, Scalar)
+from dask.dataframe.core import (repartition_divisions, _loc, aca, _concat,
+                                 _Frame, Scalar)
 from dask.dataframe.utils import eq, make_meta
 
 
@@ -1410,11 +1410,6 @@ def test_deterministic_reduction_names():
     assert a.x.min()._name == a.x.min()._name
     assert a.x.max()._name == a.x.max()._name
     assert a.x.count()._name == a.x.count()._name
-    # Test reduction without token string
-    assert (sorted(reduction(a.x, len, np.sum).dask) !=
-            sorted(reduction(a.x, np.sum, np.sum).dask))
-    assert (sorted(reduction(a.x, len, np.sum).dask) ==
-            sorted(reduction(a.x, len, np.sum).dask))
 
 
 def test_deterministic_apply_concat_apply_names():
@@ -1427,6 +1422,7 @@ def test_deterministic_apply_concat_apply_names():
             sorted(a.x.drop_duplicates().dask))
     assert (sorted(a.groupby('x').y.mean().dask) ==
             sorted(a.groupby('x').y.mean().dask))
+
     # Test aca without passing in token string
     f = lambda a: a.nlargest(5)
     f2 = lambda a: a.nlargest(3)
@@ -1435,6 +1431,32 @@ def test_deterministic_apply_concat_apply_names():
     assert (sorted(aca(a.x, f, f, a.x._meta).dask) ==
             sorted(aca(a.x, f, f, a.x._meta).dask))
 
+    # Test aca with keywords
+    def chunk(x, c_key=0, both_key=0):
+        return x.sum() + c_key + both_key
+
+    def agg(x, a_key=0, both_key=0):
+        return pd.Series(x).sum() + a_key + both_key
+
+    c_key = 2
+    a_key = 3
+    both_key = 4
+
+    res = aca(a.x, chunk=chunk, aggregate=agg, chunk_kwargs={'c_key': c_key},
+              aggregate_kwargs={'a_key': a_key}, both_key=both_key)
+    assert (sorted(res.dask) ==
+            sorted(aca(a.x, chunk=chunk, aggregate=agg,
+                       chunk_kwargs={'c_key': c_key},
+                       aggregate_kwargs={'a_key': a_key},
+                       both_key=both_key).dask))
+    assert (sorted(res.dask) !=
+            sorted(aca(a.x, chunk=chunk, aggregate=agg,
+                       chunk_kwargs={'c_key': c_key},
+                       aggregate_kwargs={'a_key': a_key},
+                       both_key=0).dask))
+
+    assert eq(res, df.x.sum() + 2*(c_key + both_key) + a_key + both_key)
+
 
 def test_aca_meta_infer():
     df = pd.DataFrame({'x': [1, 2, 3, 4],
@@ -1459,6 +1481,38 @@ def agg(x):
     assert res.compute() == df.x.sum()
 
 
+def test_reduction_method():
+    df = pd.DataFrame({'x': range(50), 'y': range(50, 100)})
+    ddf = dd.from_pandas(df, npartitions=4)
+
+    chunk = lambda x, val=0: (x >= val).sum()
+    agg = lambda x: x.sum()
+
+    # Output of chunk is a scalar
+    res = ddf.x.reduction(chunk, aggregate=agg)
+    assert eq(res, df.x.count())
+
+    # Output of chunk is a series
+    res = ddf.reduction(chunk, aggregate=agg)
+    assert res._name == ddf.reduction(chunk, aggregate=agg)._name
+    assert eq(res, df.count())
+
+    # Test with keywords
+    res2 = ddf.reduction(chunk, aggregate=agg, chunk_kwargs={'val': 25})
+    res2._name == ddf.reduction(chunk, aggregate=agg,
+                                chunk_kwargs={'val': 25})._name
+    assert res2._name != res._name
+    assert eq(res2, (df >= 25).sum())
+
+    # Output of chunk is a dataframe
+    def sum_and_count(x):
+        return pd.DataFrame({'sum': x.sum(), 'count': x.count()})
+    res = ddf.reduction(sum_and_count,
+                        aggregate=lambda x: x.groupby(level=0).sum())
+
+    assert eq(res, pd.DataFrame({'sum': df.sum(), 'count': df.count()}))
+
+
 def test_gh_517():
     arr = np.random.randn(100, 2)
     df = pd.DataFrame(arr, columns=['a', 'b'])