Repository: dask/dask

Issue:
regression in a custom groupby.agg when split_out is not None
<!-- 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**:

It seems that the changes in #8468 broke custom aggregation over in `dask_geopandas`, where we implement our own function to be able to aggregate geometries. When specifying `split_out` in the `.agg` method, it currently raises an error. With `split_out=None`, it is all fine, so I think that this is an issue on the dask side, not on dask_geopandas side.

**What you expected to happen**:

I expect that the aggregation will work no matter if I specify `split_out` or not.

**Minimal Complete Verifiable Example**:

This example requires `geopandas` and `dask_geopandas`, I wasn't able to craft one with pure dask.
```python
import geopandas
import dask_geopandas

world = geopandas.read_file(geopandas.datasets.get_path("naturalearth_lowres"))
ddf = dask_geopandas.from_geopandas(world, npartitions=4)


def union(block):
    merged_geom = block.unary_union
    return merged_geom

merge_geometries = dd.Aggregation(
    "merge_geometries", lambda s: s.agg(union), lambda s0: s0.agg(union)
)
data_agg = {col: aggfunc for col in ddf.columns.drop([ddf.geometry.name])}
data_agg[ddf.geometry.name] = merge_geometries

# this works
# aggregated = ddf.groupby(by="continent").agg(data_agg)

# this doesn't
aggregated = ddf.groupby(by="continent").agg(data_agg, split_out=2)
```

Traceback:
<details>
```py
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
~/mambaforge/envs/geo_dev/lib/python3.9/site-packages/pandas/core/indexes/base.py in get_loc(self, key, method, tolerance)
   3360             try:
-> 3361                 return self._engine.get_loc(casted_key)
   3362             except KeyError as err:

~/mambaforge/envs/geo_dev/lib/python3.9/site-packages/pandas/_libs/index.pyx in pandas._libs.index.IndexEngine.get_loc()

~/mambaforge/envs/geo_dev/lib/python3.9/site-packages/pandas/_libs/index.pyx in pandas._libs.index.IndexEngine.get_loc()

pandas/_libs/hashtable_class_helper.pxi in pandas._libs.hashtable.PyObjectHashTable.get_item()

pandas/_libs/hashtable_class_helper.pxi in pandas._libs.hashtable.PyObjectHashTable.get_item()

KeyError: 'geometry'

The above exception was the direct cause of the following exception:

KeyError                                  Traceback (most recent call last)
~/Git/geopandas/geopandas/geodataframe.py in set_geometry(self, col, drop, inplace, crs)
    310             try:
--> 311                 level = frame[col]
    312             except KeyError:

~/Git/geopandas/geopandas/geodataframe.py in __getitem__(self, key)
   1349         """
-> 1350         result = super().__getitem__(key)
   1351         geo_col = self._geometry_column_name

~/mambaforge/envs/geo_dev/lib/python3.9/site-packages/pandas/core/frame.py in __getitem__(self, key)
   3457                 return self._getitem_multilevel(key)
-> 3458             indexer = self.columns.get_loc(key)
   3459             if is_integer(indexer):

~/mambaforge/envs/geo_dev/lib/python3.9/site-packages/pandas/core/indexes/base.py in get_loc(self, key, method, tolerance)
   3362             except KeyError as err:
-> 3363                 raise KeyError(key) from err
   3364 

KeyError: 'geometry'

During handling of the above exception, another exception occurred:

ValueError                                Traceback (most recent call last)
~/mambaforge/envs/geo_dev/lib/python3.9/site-packages/dask/dataframe/utils.py in raise_on_meta_error(funcname, udf)
    174     try:
--> 175         yield
    176     except Exception as e:

~/mambaforge/envs/geo_dev/lib/python3.9/site-packages/dask/dataframe/core.py in _emulate(func, udf, *args, **kwargs)
   5965     with raise_on_meta_error(funcname(func), udf=udf):
-> 5966         return func(*_extract_meta(args, True), **_extract_meta(kwargs, True))
   5967 

~/mambaforge/envs/geo_dev/lib/python3.9/site-packages/dask/dataframe/core.py in _extract_meta(x, nonempty)
   5945     elif isinstance(x, tuple):
-> 5946         return tuple(_extract_meta(_x, nonempty) for _x in x)
   5947     elif isinstance(x, dict):

~/mambaforge/envs/geo_dev/lib/python3.9/site-packages/dask/dataframe/core.py in <genexpr>(.0)
   5945     elif isinstance(x, tuple):
-> 5946         return tuple(_extract_meta(_x, nonempty) for _x in x)
   5947     elif isinstance(x, dict):

~/mambaforge/envs/geo_dev/lib/python3.9/site-packages/dask/dataframe/core.py in _extract_meta(x, nonempty)
   5941     if isinstance(x, (Scalar, _Frame)):
-> 5942         return x._meta_nonempty if nonempty else x._meta
   5943     elif isinstance(x, list):

~/mambaforge/envs/geo_dev/lib/python3.9/site-packages/dask/dataframe/core.py in _meta_nonempty(self)
    430         """A non-empty version of `_meta` with fake data."""
--> 431         return meta_nonempty(self._meta)
    432 

~/mambaforge/envs/geo_dev/lib/python3.9/site-packages/dask/utils.py in __call__(self, arg, *args, **kwargs)
    623         meth = self.dispatch(type(arg))
--> 624         return meth(arg, *args, **kwargs)
    625 

~/Git/dask-geopandas/dask_geopandas/backends.py in _nonempty_geodataframe(x)
     59     df = meta_nonempty_dataframe(x)
---> 60     return geopandas.GeoDataFrame(df, geometry=x._geometry_column_name, crs=x.crs)
     61 

~/Git/geopandas/geopandas/geodataframe.py in __init__(self, data, geometry, crs, *args, **kwargs)
    203                 # TODO: raise error in 0.9 or 0.10.
--> 204             self.set_geometry(geometry, inplace=True)
    205 

~/Git/geopandas/geopandas/geodataframe.py in set_geometry(self, col, drop, inplace, crs)
    312             except KeyError:
--> 313                 raise ValueError("Unknown column %s" % col)
    314             except Exception:

ValueError: Unknown column geometry

The above exception was the direct cause of the following exception:

ValueError                                Traceback (most recent call last)
/var/folders/2f/fhks6w_d0k556plcv3rfmshw0000gn/T/ipykernel_44916/2101873994.py in <module>
     12 data_agg = {col: aggfunc for col in ddf.columns.drop([ddf.geometry.name])}
     13 data_agg[ddf.geometry.name] = merge_geometries
---> 14 aggregated = ddf.groupby(by="continent").agg(data_agg, split_out=2)

~/mambaforge/envs/geo_dev/lib/python3.9/site-packages/dask/dataframe/groupby.py in agg(self, arg, split_every, split_out)
   2006     @derived_from(pd.core.groupby.DataFrameGroupBy)
   2007     def agg(self, arg, split_every=None, split_out=1):
-> 2008         return self.aggregate(arg, split_every=split_every, split_out=split_out)
   2009 
   2010 

~/mambaforge/envs/geo_dev/lib/python3.9/site-packages/dask/dataframe/groupby.py in aggregate(self, arg, split_every, split_out)
   2002             return self.size()
   2003 
-> 2004         return super().aggregate(arg, split_every=split_every, split_out=split_out)
   2005 
   2006     @derived_from(pd.core.groupby.DataFrameGroupBy)

~/mambaforge/envs/geo_dev/lib/python3.9/site-packages/dask/dataframe/groupby.py in aggregate(self, arg, split_every, split_out)
   1619             )
   1620 
-> 1621         return aca(
   1622             chunk_args,
   1623             chunk=_groupby_apply_funcs,

~/mambaforge/envs/geo_dev/lib/python3.9/site-packages/dask/dataframe/core.py in apply_concat_apply(args, chunk, aggregate, combine, meta, token, chunk_kwargs, aggregate_kwargs, combine_kwargs, split_every, split_out, split_out_setup, split_out_setup_kwargs, sort, ignore_index, **kwargs)
   5878     # Blockwise Split Layer
   5879     if split_out and split_out > 1:
-> 5880         chunked = chunked.map_partitions(
   5881             hash_shard,
   5882             split_out,

~/mambaforge/envs/geo_dev/lib/python3.9/site-packages/dask/dataframe/core.py in map_partitions(self, func, *args, **kwargs)
    769         None as the division.
    770         """
--> 771         return map_partitions(func, self, *args, **kwargs)
    772 
    773     @insert_meta_param_description(pad=12)

~/mambaforge/envs/geo_dev/lib/python3.9/site-packages/dask/dataframe/core.py in map_partitions(func, meta, enforce_metadata, transform_divisions, align_dataframes, *args, **kwargs)
   6033         # Use non-normalized kwargs here, as we want the real values (not
   6034         # delayed values)
-> 6035         meta = _emulate(func, *args, udf=True, **kwargs)
   6036     else:
   6037         meta = make_meta(meta, index=meta_index, parent_meta=parent_meta)

~/mambaforge/envs/geo_dev/lib/python3.9/site-packages/dask/dataframe/core.py in _emulate(func, udf, *args, **kwargs)
   5964     """
   5965     with raise_on_meta_error(funcname(func), udf=udf):
-> 5966         return func(*_extract_meta(args, True), **_extract_meta(kwargs, True))
   5967 
   5968 

~/mambaforge/envs/geo_dev/lib/python3.9/contextlib.py in __exit__(self, typ, value, traceback)
    135                 value = typ()
    136             try:
--> 137                 self.gen.throw(typ, value, traceback)
    138             except StopIteration as exc:
    139                 # Suppress StopIteration *unless* it's the same exception that

~/mambaforge/envs/geo_dev/lib/python3.9/site-packages/dask/dataframe/utils.py in raise_on_meta_error(funcname, udf)
    194         )
    195         msg = msg.format(f" in `{funcname}`" if funcname else "", repr(e), tb)
--> 196         raise ValueError(msg) from e
    197 
    198 

ValueError: Metadata inference failed in `hash_shard`.

You have supplied a custom function and Dask is unable to 
determine the type of output that that function returns. 

To resolve this please provide a meta= keyword.
The docstring of the Dask function you ran should have more information.

Original error is below:
------------------------
ValueError('Unknown column geometry')

Traceback:
---------
  File "/Users/martin/mambaforge/envs/geo_dev/lib/python3.9/site-packages/dask/dataframe/utils.py", line 175, in raise_on_meta_error
    yield
  File "/Users/martin/mambaforge/envs/geo_dev/lib/python3.9/site-packages/dask/dataframe/core.py", line 5966, in _emulate
    return func(*_extract_meta(args, True), **_extract_meta(kwargs, True))
  File "/Users/martin/mambaforge/envs/geo_dev/lib/python3.9/site-packages/dask/dataframe/core.py", line 5946, in _extract_meta
    return tuple(_extract_meta(_x, nonempty) for _x in x)
  File "/Users/martin/mambaforge/envs/geo_dev/lib/python3.9/site-packages/dask/dataframe/core.py", line 5946, in <genexpr>
    return tuple(_extract_meta(_x, nonempty) for _x in x)
  File "/Users/martin/mambaforge/envs/geo_dev/lib/python3.9/site-packages/dask/dataframe/core.py", line 5942, in _extract_meta
    return x._meta_nonempty if nonempty else x._meta
  File "/Users/martin/mambaforge/envs/geo_dev/lib/python3.9/site-packages/dask/dataframe/core.py", line 431, in _meta_nonempty
    return meta_nonempty(self._meta)
  File "/Users/martin/mambaforge/envs/geo_dev/lib/python3.9/site-packages/dask/utils.py", line 624, in __call__
    return meth(arg, *args, **kwargs)
  File "/Users/martin/Git/dask-geopandas/dask_geopandas/backends.py", line 60, in _nonempty_geodataframe
    return geopandas.GeoDataFrame(df, geometry=x._geometry_column_name, crs=x.crs)
  File "/Users/martin/Git/geopandas/geopandas/geodataframe.py", line 204, in __init__
    self.set_geometry(geometry, inplace=True)
  File "/Users/martin/Git/geopandas/geopandas/geodataframe.py", line 313, in set_geometry
    raise ValueError("Unknown column %s" % col)
```
</details>

**Anything else we need to know?**:

**Environment**:

- Dask version: '2022.01.0+17.g08eee6d6'
- Python version: 3.9.7
- Operating System: macOS, ubuntu
- Install method (conda, pip, source): pip + git

<!-- If you are reporting an issue such as scale stability, cluster deadlock.
Please provide a cluster dump state with this issue, by running client.dump_cluster_state()

https://distributed.dask.org/en/stable/api.html?highlight=dump_cluster_state#distributed.Client.dump_cluster_state

-->

<details>
<summary>Cluster Dump State:</summary>

</details>


Discussion:
It's probably not a regression in dask, but an existing issue with the interaction with geopandas.GeoDataFrame being a special sublcass. 

What changed in #8468, is that `apply_concat_apply` starts to use `map_partitions`. And with `split_out` specified, it calls `map_partitions` a second time, which I think is what triggers the direct error. But I think it's actually already caused by the first call of `map_partitions`. 
In our case of the `dissolve` method, this is basically a wrapper around `groupby`. And then, the function being mapped is `_groupby_apply_funcs`, which has this part in the code:

https://github.com/dask/dask/blob/08eee6d655459e9b8494d67957b6fe908e1b8c67/dask/dataframe/groupby.py#L929-L930

where `result` is a dictionary of columns, but with renamed column names.

A very simple example that reproduces the underlying issue (without involving groupby / apply_concat_apply):

```python
import geopandas
import dask_geopandas

df = geopandas.read_file(geopandas.datasets.get_path("naturalearth_lowres"))
ddf = dask_geopandas.from_geopandas(df, npartitions=4)

def func(df):
    result = {}
    for col in df.columns:
        result["prefix-" + col] = df[col]
    # naively recreating the class, without accounting for the renamed geometry column
    return df.__class__(result)


>>> res = ddf.map_partitions(func)
>>> res
Dask GeoDataFrame Structure:
              prefix-pop_est prefix-continent prefix-name prefix-iso_a3 prefix-gdp_md_est prefix-geometry
npartitions=4                                                                                            
0                      int64           object      object        object           float64        geometry
45                       ...              ...         ...           ...               ...             ...
90                       ...              ...         ...           ...               ...             ...
135                      ...              ...         ...           ...               ...             ...
176                      ...              ...         ...           ...               ...             ...
Dask Name: func, 8 tasks

>>> res._meta
Empty GeoDataFrame
Columns: [prefix-pop_est, prefix-continent, prefix-name, prefix-iso_a3, prefix-gdp_md_est, prefix-geometry]
Index: []

>>> res._meta._geometry_column_name
'geometry'
```

So we get a dask_geopandas GeoDataFrame where the `_meta` is an "invalid" GeoDataFrame, because the geometry column name is not correctly set. 

Since dask cannot know that it should handle GeoDataFrame construction in a special way inside `_groupby_apply_funcs`, this is probably something to discuss on the geopandas side on how we can deal with this.
Thanks for investigating this @jorisvandenbossche! Feel free to ping me if we need changes on the Dask side
Maybe a general question about the groupby logic in dask: why is it needed to rename the columns (with prefix about the task and hash suffix) in the intermediary objects? 
(it's this renaming of columns that causes issues for geopandas)
The renaming seems to happen at

https://github.com/dask/dask/blob/08eee6d655459e9b8494d67957b6fe908e1b8c67/dask/dataframe/groupby.py#L574-L575

If I naively change that to return the original `column`, the example in the top post works again (but there are also a whole bunch of groupby tests that fail). Now, I suppose there is a good reason for this renaming of the aggregated columns, but I don't directly understand it. 
We have discussed this issue with @jorisvandenbossche and it is unlikely that we will be able to fix it on geopandas or dask-geopandas side. 

`geopandas.GeoDataFrame` defines a `.geometry` property that takes care of all geospatial operations (interfacing shapely and pygeos). `.geometry` always points to one of the columns of the `GeoDataFrame` based on its name ("geometry" by default but can be any). When dask renames the columns in groupby, this link is broken and any geospatial operation called within groupby (like the `unary_union` above) fails because geopandas doesn't know which column is the active geometry anymore as the name it expects does no longer exist in the GeoDataFrame. Since there can be multiple geometry columns (GeoSeries), we cannot infer this active geometry from the dtype and using positional index also does not seem to be feasible.

I am not certain what has changed in dask now, but it seems that until #8468 the renaming either did not happen or happened in a way that did not break this link column<->`.geometry`. But I think that we'll need your help on the dask side to fix this issue.
Is this a regression we need to fix or revert before releasing https://github.com/dask/community/issues/216? cc @jrbourbeau

@martinfleis @jorisvandenbossche do you have a sense of what the fix would be? Not renaming the columns in groupby seems like one idea—I haven't checked yet to understand why we do that. But I get the sense that the more robust fix would be to use something more advanced than: https://github.com/dask/dask/blob/08eee6d655459e9b8494d67957b6fe908e1b8c67/dask/dataframe/groupby.py#L929-L930

Even if we stopped renaming the columns, if a user was using a geometry column that wasn't named `geometry`, wouldn't this still break? It seems to me like we need some interface to get extra arguments into the dataframe constructor (or to specify a custom constructor function, which you could `functools.partial` as needed).
> Is this a regression we need to fix or revert before releasing dask/community#216?

It breaks `dissolve()` in dask-geopandas when `split_out` is specified but we haven't released this yet (will do in a week or two, hopefully). So if there's a perspective that this will get fixed in on of the upcoming versions there's no need to rush for 2022.01.1.

> Even if we stopped renaming the columns, if a user was using a geometry column that wasn't named geometry, wouldn't this still break?

It probably would. When the GeoDataFrame is constructed from a dict as in this case, geopandas either checks for "geometry" or needs the geometry column to be specified. I think that we would need to pass a custom constructor that would get the information on the geometry column from `_meta`? I am not really sure, my insight into dask internals is quite limited. Maybe @jorisvandenbossche got some ideas.
Hmm - I am having a hard time figuring out how #8468 (or setting split_out) ended up changing behavior like this. It really seems like `_groupby_apply_funcs` has always initialized a new dataframe-like object with modified column names.  You are correct that we do not need to rename columns in all cases, but we certainly do need to create new names if we are performing more than one type of aggregation on the same column (pretty common).  For example, we may want the "min" and "max" of the same column, or we may need the count and the sum to calculate mean in the final output.

I'm really hoping to find exactly what changed in #8468, but it also seems like the assumption that this `geometry` parameter is always defined may be a bit tricky in the case of a ACA.  Is it possible to point this property to a fallback function that raises an error if the user tries to do anything that requires it? I realize you would typically want to let the user know right away when they initialize an invalid object, so no worries if this seems unattractive.  What would be the ideal approach in the case that you are performing multiple aggregations on the "geometry" column here? Would you want both to be referenced by `geometry`?
@rjzamora It had to be some commit merged on 22/1 and #8468 seems to be the only candidate.

> Is it possible to point this property to a fallback function that raises an error if the user tries to do anything that requires it?

I am not sure what do you mean but we need to use it, we can't just raise an informative error. The whole point of dask-geopandas is that is enables geospatial operations and the needs to include aggregations. The `dissolve` example we are talking about here is one of the most common use cases.

The invalid object is not created by user, it is created by dask and user has no way of influencing it as far as I can tell. 

> What would be the ideal approach in the case that you are performing multiple aggregations on the "geometry" column here? Would you want both to be referenced by `geometry`?

The actual aggregation of the column is not necessarily problematic. If you get a GeoSeries, you no longer have the issue with unset `geometry`, and the geometric operation just works. So they can be called whatever you need them to be called and it should not affect the aggregation. In the example above, the error does not come from calling the geometric `unary_union` but from recreating meta if I understand the error correctly. [This](https://github.com/geopandas/dask-geopandas/blob/791751d3f3110fa8df09dfcbbad2ed643f6a79b0/dask_geopandas/backends.py#L57-L60) is where the issue is raised:

```py
@meta_nonempty.register(geopandas.GeoDataFrame)
def _nonempty_geodataframe(x):
    df = meta_nonempty_dataframe(x)
    return geopandas.GeoDataFrame(df, geometry=x._geometry_column_name, crs=x.crs)
```

Once Dask calls this after it renames the columns, `x._geometry_column_name` which is a string, usually `"geometry"` cannot find a column of the same name and raises a ValueError.
>if I understand the error correctly. This is where the issue is raised:

 Ah - I was completely focused on the `return df.__class__(result)` discussion above.  However, I now realize that `GeoDataFrame(result)` will work fine when `result` is a dict without a “geometry” field.  All of my questions/suggestions above were assuming that something like that will fail.  The only time that “geometry” constraint shows up is when you try to create a `GeoDataFrame` metadata object.  Now it makes perfect sense why moving to `map_partitions` in #8468 is changing behavior.
It seems like the problem may only show up when a "non-empty" metadata object needs to be generated by geopandas. I experimented with some quick things in ACA, and it seems like [this branch](https://github.com/rjzamora/dask/tree/empty-meta-emulation) (which avoids the need to non-empty metadata in the`has_shard` operation relevent for `split_out>1`) **may** resolve this issue.
@rjzamora `dask-geopandas` CI passes with this branch, so it may be a solution. Thanks for looking into that!
> Hmm - I am having a hard time figuring out how #8468 (or setting split_out) ended up changing behavior like this. It really seems like `_groupby_apply_funcs` has always initialized a new dataframe-like object with modified column names.

Given your last comments it might now be clear to you, but just to reiterate in general, my understanding is that:

* The groupby implementation renames columns and recreates a GeoDataFrame (the `return df.__class__(result)` where `result` is a dict with renamed keys). This in itself is not a problem (geopandas will let that pass fine, although you get an "invalid" geodataframe, that's not a direct problem because the columns get renamed back to the original ones before the end result is passed to the user). But, it causes a problem in `_meta_nonempty` if this is called on the invalid GeoDataFrame (so if it is called *before* the columns are renamed back to the original). 
* #8468 started to use `map_partitions` in the `aca` implementation that is used in `groupby.agg(..)`, and `map_partitions` calls `_meta_nonempty` (though `_emulate`: https://github.com/dask/dask/blob/04ade86291f0f8e7eb6ddd36e06dfb722252345b/dask/dataframe/core.py#L6041-L6044, which always calls `_extract_meta` with `nonempty=True`). 
  So indirectly, [#8468](https://github.com/dask/dask/pull/8468) started to call `_meta_nonempty`.
* With the default `groupby.agg(..)`/`aca`, `map_partitions` gets only called once, and at that point the `_meta_nonempty` gets constructed _before_ the columns are renamed in the `_groupby_apply_funcs`. 
  But if using `split_out`, `map_partitions` gets called a second time, and then `_meta_nonemtpy` gets constructed again, and now with the renamed columns, giving the error in constructing the nonempty GeoDataFrame.

---

So indeed if we could avoid the use of `_meta_nonempty` in `map_partitions`, that would solve this issue. 

While writing this up, I also thought of a potential solution on the dask-geopandas side in being a bit more flexible in creating the meta-nonempty:

```diff
--- a/dask_geopandas/backends.py
+++ b/dask_geopandas/backends.py
@@ -56,8 +56,12 @@ def _nonempty_geoseries(x, idx=None):
 
 @meta_nonempty.register(geopandas.GeoDataFrame)
 def _nonempty_geodataframe(x):
     df = meta_nonempty_dataframe(x)
-    return geopandas.GeoDataFrame(df, geometry=x._geometry_column_name, crs=x.crs)
+    try:
+        return geopandas.GeoDataFrame(df, geometry=x._geometry_column_name, crs=x.crs)
+    except:
+        return geopandas.GeoDataFrame(df)
```

This also "solves" the issue for this specific case. I am only not sure if it won't give problems in other cases. 
The reason this fixes the issue here for groupby, is because this will not error when creating the `_meta_nonempty` the second time with the renamed columns, and because the columns are renamed back to the original ones before returning the result to the user, the end result is a proper GeoDataFrame (this will also only work with the default "geometry" column name). 



@jorisvandenbossche I think that the solution avoiding call to `meta_nonempty` is better here because, as you pointed out, this would work only for the geometry column named `"geometry"`. It could allow a workaround but if we can make it generalisable for any name, that would be optimal. 

Resolving patch:
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -5878,6 +5878,12 @@ def apply_concat_apply(
         chunk,
         *args,
         token=chunk_name,
+        # NOTE: We are NOT setting the correct
+        # `meta` here on purpose. We are using
+        # `map_partitions` as a convenient way
+        # to build a `Blockwise` layer, and need
+        # to avoid the metadata emulation step.
+        meta=dfs[0],
         enforce_metadata=False,
         transform_divisions=False,
         align_dataframes=False,
@@ -5893,6 +5899,12 @@ def apply_concat_apply(
             split_out_setup_kwargs,
             ignore_index,
             token="split-%s" % token_key,
+            # NOTE: We are NOT setting the correct
+            # `meta` here on purpose. We are using
+            # `map_partitions` as a convenient way
+            # to build a `Blockwise` layer, and need
+            # to avoid the metadata emulation step.
+            meta=dfs[0],
             enforce_metadata=False,
             transform_divisions=False,
             align_dataframes=False,


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
@@ -5052,3 +5052,34 @@ def test_index_is_monotonic_dt64():
     s_2 = pd.Series(1, index=list(reversed(s)))
     ds_2 = dd.from_pandas(s_2, npartitions=5, sort=False)
     assert_eq(s_2.index.is_monotonic_decreasing, ds_2.index.is_monotonic_decreasing)
+
+
+def test_custom_map_reduce():
+    # Make sure custom map-reduce workflows can use
+    # the universal ACA code path with metadata
+    # that is not DataFrame-like.
+    # See: https://github.com/dask/dask/issues/8636
+
+    df = pd.DataFrame(columns=["a"], data=[[2], [4], [8]], index=[0, 1, 2])
+    ddf = dd.from_pandas(df, npartitions=2)
+
+    def map_fn(x):
+        return {"x": x, "y": x}
+
+    def reduce_fn(series):
+        merged = None
+        for mapped in series:
+            if merged is None:
+                merged = mapped.copy()
+            else:
+                merged["x"] += mapped["x"]
+                merged["y"] *= mapped["y"]
+        return merged
+
+    result = (
+        ddf["a"]
+        .map(map_fn, meta=("data", "object"))
+        .reduction(reduce_fn, aggregate=reduce_fn, meta=("data", "object"))
+        .compute()[0]
+    )
+    assert result == {"x": 14, "y": 64}