Repository: dask/dask

Issue:
Series Aggregation called Multiple Times, Must Support Recursive Nature
<!-- Please include a self-contained copy-pastable example that generates the issue if possible.

Please be concise with code posted. See guidelines below on how to provide a good bug report:

- Craft Minimal Bug Reports http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports
- Minimal Complete Verifiable Examples https://stackoverflow.com/help/mcve

Bug reports that follow these guidelines are easier to diagnose, and so are often handled much more quickly.
-->

**What happened**:

With the latest Dask version 2022.02.0, whenever a Series is aggregated, the aggregator is called multiple times in a recursive nature. This was not the case before 2022.02.0. This recursive nature means that the aggregator function must return values that it should be able to later accept again.

**What you expected to happen**:

In Dask 2022.01.0 and before, the aggregator did not need to handle this recursive nature. It could return a value which it did not need to accept in a future call.

**Minimal Complete Verifiable Example**:

```python
import dask.dataframe as dd
from dask import distributed
import pandas as pd

class SomeObject():
    def __init__(self, val):
        self.val = val
        # Print for clarity what is going on
        print(self.val)

    def __str__(self):
        return f"Value: {self.val}"

def _id(x):
    return x

def aggregate_something(df_mps):
    return SomeObject(df_mps.sum().sum())

def main():
    # If you comment out this line, things work fine
    client = distributed.Client()

    N = 170
    series = pd.Series(data=[1] * N, index=range(2, N + 2))
    dask_series = dd.from_pandas(series, npartitions=100)

    res = dask_series.reduction(
        chunk=_id,
        aggregate=aggregate_something,
        split_every=False,
        token="commit-dataset",
        meta=object,
    )

    print(res.compute())

if __name__ == "__main__":
    main()
```

The *broken* output with Dask 2022.02.0:

<details>
<summary>2022.02.0 Output</summary>

```bash
272000.0
272000.0
68000.0
272000.0
272000.0
272000.0
272000.0
distributed.worker - WARNING - Compute Failed
Function:  execute_task
args:      ((<function pipe at 0x7f90e95f7670>, [<__main__.SomeObject object at 0x7f90e363a220>, <__main__.SomeObject object at 0x7f90e363a280>, <__main__.SomeObject object at 0x7f90269cb9a0>, <__main__.SomeObject object at 0x7f90e363ae20>, <__main__.SomeObject object at 0x7f90122909a0>, <__main__.SomeObject object at 0x7f90e363aa30>, <__main__.SomeObject object at 0x7f9011e5f190>], functools.partial(<function _concat at 0x7f9056c23430>, ignore_index=False), functools.partial(<function _reduction_aggregate at 0x7f9056c4ab80>, aca_aggregate=<function aggregate_something at 0x7f9011d5f3a0>)))
kwargs:    {}
Exception: 'TypeError("unsupported operand type(s) for +: \'SomeObject\' and \'SomeObject\'")'

Traceback (most recent call last):
  File "/home/X931300/debugging/playground_dask/dask_kartothek_reduction_failure.py", line 37, in <module>
    main()
  File "/home/X931300/debugging/playground_dask/dask_kartothek_reduction_failure.py", line 34, in main
    print(res.compute(scheduler="distributed"))
  File "/home/X931300/debugging/dask/dask/base.py", line 291, in compute
    (result,) = compute(self, traverse=False, **kwargs)
  File "/home/X931300/debugging/dask/dask/base.py", line 574, in compute
    results = schedule(dsk, keys, **kwargs)
  File "/tmp/debug/lib/python3.9/site-packages/distributed/client.py", line 3010, in get
    results = self.gather(packed, asynchronous=asynchronous, direct=direct)
  File "/tmp/debug/lib/python3.9/site-packages/distributed/client.py", line 2162, in gather
    return self.sync(
  File "/tmp/debug/lib/python3.9/site-packages/distributed/utils.py", line 311, in sync
    return sync(
  File "/tmp/debug/lib/python3.9/site-packages/distributed/utils.py", line 378, in sync
    raise exc.with_traceback(tb)
  File "/tmp/debug/lib/python3.9/site-packages/distributed/utils.py", line 351, in f
    result = yield future
  File "/tmp/debug/lib/python3.9/site-packages/tornado/gen.py", line 762, in run
    value = future.result()
  File "/tmp/debug/lib/python3.9/site-packages/distributed/client.py", line 2025, in _gather
    raise exception.with_traceback(traceback)
  File "/tmp/debug/lib/python3.9/site-packages/toolz/functoolz.py", line 630, in pipe
    data = func(data)
  File "/home/X931300/debugging/dask/dask/dataframe/core.py", line 7136, in _reduction_aggregate
    return aca_aggregate(x, **kwargs)
  File "/home/X931300/debugging/playground_dask/dask_kartothek_reduction_failure.py", line 18, in aggregate_something
    return SomeObject(df_mps.sum().sum())
  File "/tmp/debug/lib/python3.9/site-packages/pandas/core/generic.py", line 10708, in sum
    return NDFrame.sum(
  File "/tmp/debug/lib/python3.9/site-packages/pandas/core/generic.py", line 10446, in sum
    return self._min_count_stat_function(
  File "/tmp/debug/lib/python3.9/site-packages/pandas/core/generic.py", line 10428, in _min_count_stat_function
    return self._reduce(
  File "/tmp/debug/lib/python3.9/site-packages/pandas/core/series.py", line 4392, in _reduce
    return op(delegate, skipna=skipna, **kwds)
  File "/tmp/debug/lib/python3.9/site-packages/pandas/core/nanops.py", line 94, in _f
    return f(*args, **kwargs)
  File "/tmp/debug/lib/python3.9/site-packages/pandas/core/nanops.py", line 411, in new_func
    result = func(values, axis=axis, skipna=skipna, mask=mask, **kwargs)
  File "/tmp/debug/lib/python3.9/site-packages/pandas/core/nanops.py", line 591, in nansum
    the_sum = values.sum(axis, dtype=dtype_sum)
  File "/tmp/debug/lib/python3.9/site-packages/numpy/core/_methods.py", line 48, in _sum
    return umr_sum(a, axis, dtype, out, keepdims, initial, where)
TypeError: unsupported operand type(s) for +: 'SomeObject' and 'SomeObject'
```
</details>

The *previous* output with Dask 2022.01.0:

<details>
<summary>2022.01.0 Output</summary>

```bash
1700000.0
Value: 1700000.0
```

</details>

**Anything else we need to know?**:

This bug presented itself in the [Kartothek v4.0.3](https://github.com/JDASoftwareGroup/kartothek/tree/db840b52f28746f128e6aadee2103e6021db0e1a) library, specifically in a Dask Series reduction operation to update Kartothek's metadata [here](https://github.com/JDASoftwareGroup/kartothek/blob/v4.0.3/kartothek/io/dask/dataframe.py#L449). The reduction function it uses `_commit_update_from_reduction` ultimately ends up [here](_commit_update_from_reduction) which takes in a list of lists of [MetaPartition](https://github.com/JDASoftwareGroup/kartothek/blob/v4.0.3/kartothek/io_components/metapartition.py#L203) objects and returns a [DatasetMetadata](https://github.com/JDASoftwareGroup/kartothek/blob/4d20267393d44a31a895077f81f52f1f730557c4/kartothek/core/dataset.py#L502) object. The important thing to note is that this function is not recursive-ready -- it cannot be recursed on the value it outputs, the new requirement of a Dask Series reduction function.

The PR in Dask which introduces this altered behavior is: https://github.com/dask/dask/pull/8468 . All commit before that PR did not have this issue.

**Environment**:

- Dask version: 2022.02.0
- Python version: Python 3.9.10
- Operating System: Ubuntu 20.04.3 LTS
- Install method (conda, pip, source): source



Discussion:
Thanks for the reproducible example, I've updated it slightly to run with less data. I also noted that this works fine locally and only fails with the distributed scheduler, indicating this likely has something to do with the HighLevelGraph logic for generating the graph on the scheduler.

cc @rjzamora, since #8468 lead to this issue.
Your intuition sounds correct to me @jcrist - The good news is that #8672 fixes the problem. However, I will see if I can find what is causing the issue now in case that PR takes a while to get merged.

(**SECOND EDIT**: The Layer refactor indeed "fixes" the bug. My earlier hesitation was because a "simplified" version of the repro with `split_out=2` was failing for me.  However, anything other than `split_out=False` will always fail here since `aggregate_something` doesn't support tree aggregation anyway - My mistake.)

Resolving patch:
diff --git a/dask/layers.py b/dask/layers.py
--- a/dask/layers.py
+++ b/dask/layers.py
@@ -1515,6 +1515,7 @@ def __dask_distributed_pack__(self, *args, **kwargs):
             "concat_func": _concat_func,
             "tree_node_func": _tree_node_func,
             "finalize_func": _finalize_func,
+            "split_every": self.split_every,
             "split_out": self.split_out,
             "output_partitions": self.output_partitions,
             "tree_node_name": self.tree_node_name,


Test patch:
diff --git a/dask/tests/test_distributed.py b/dask/tests/test_distributed.py
--- a/dask/tests/test_distributed.py
+++ b/dask/tests/test_distributed.py
@@ -698,3 +698,28 @@ async def test_to_sql_engine_kwargs(c, s, a, b):
             dd.read_sql_table("test", uri, "index"),
             check_divisions=False,
         )
+
+
+@gen_cluster(client=True)
+async def test_non_recursive_df_reduce(c, s, a, b):
+    # See https://github.com/dask/dask/issues/8773
+
+    dd = pytest.importorskip("dask.dataframe")
+    pd = pytest.importorskip("pandas")
+
+    class SomeObject:
+        def __init__(self, val):
+            self.val = val
+
+    N = 170
+    series = pd.Series(data=[1] * N, index=range(2, N + 2))
+    dask_series = dd.from_pandas(series, npartitions=34)
+    result = dask_series.reduction(
+        chunk=lambda x: x,
+        aggregate=lambda x: SomeObject(x.sum().sum()),
+        split_every=False,
+        token="commit-dataset",
+        meta=object,
+    )
+
+    assert (await c.compute(result)).val == 170