DataFrame subclass lost in `groupby.agg` with `split_out` set.
- Dominant language
- Python
- Stars
- 89
- Forks
- 26
- PR merge metrics
- No merged PRs in 30d
Description
**Describe the issue**:
As part of https://github.com/geopandas/dask-geopandas/pull/285, we found that dask-expr will lose the type of a pandas DataFrame subclass in `groupby.agg` if (and only if?) the `split_out` parameter is used.
**Minimal Complete Verifiable Example**:
Given this file:
```python
# file: test.py
import dask.dataframe.backends
import pandas as pd
import dask_expr as dx
import dask.dataframe as dd
from dask.dataframe.dispatch import make_meta_dispatch, meta_nonempty
from dask.dataframe.core import get_parallel_type
import dask.dataframe.backends
dask.config.set(scheduler="single-threaded")
class MySeries(pd.Series):
@property
def _constructor(self):
return MySeries
@property
def _constructor_expanddim(self):
return MyDataFrame
class MyDataFrame(pd.DataFrame):
@property
def _constructor(self):
return MyDataFrame
@property
def _constructor_sliced(self):
return MySeries
class MyIndex(pd.Index): ...
class MyDaskSeries(dx.Series):
_partition_type = MySeries
class MyDaskDataFrame(dx.DataFrame):
_partition_type = MyDataFrame
class MyDaskIndex(dx.Index):
_partition_type = MyIndex
# Unclear if any of get_parallel_type and make_meta_dispatch are needed.
# Reproduces with or without them.
@get_parallel_type.register(MyDataFrame)
def get_parallel_type_dataframe(df):
return MyDataFrame
@get_parallel_type.register(MySeries)
def get_parallel_type_series(s):
return MyDaskSeries
@get_parallel_type.register(MyIndex)
def get_parallel_type_index(ind):
return MyDaskIndex
@make_meta_dispatch.register(MyDataFrame)
def make_meta_dataframe(df, index=None):
return df.head(0)
@make_meta_dispatch.register(MySeries)
def make_meta_series(s, index=None):
return s.head(0)
@make_meta_dispatch.register(MyIndex)
def make_meta_index(ind, index=None):
return ind[:0]
@meta_nonempty.register(MyDataFrame)
def make_meta_nonempty_dataframe(x):
return MyDataFrame(dask.dataframe.backends.meta_nonempty_dataframe(x))
df = dx.from_dict(
{"a": [1, 1, 2, 2], "b": [1, 2, 3, 4]}, npartitions=4, constructor=MyDataFrame
)
a = df.groupby("a").agg("first")
b = df.groupby("a").agg("first", split_out=2)
print("split-out=None", type(a.compute()))
print("split-out=2 ", type(b.compute()))
```
running that produces
```
$ python test.py
split-out=None
split-out=2
```
I would expect the `type` there to be `__main__.MyDataFrame` regardless of `split_out`.
**Anything else we need to know?**:
**Environment**:
```
dask 2024.4.1
dask-expr 1.0.11
```
Edit: I made one addition to the script: adding a `@meta_nonempty.register(MyDataFrame)`. I noticed that in `DecomposableGroupbyAggregation.combine` and `DecomposableGroupbyAggregation.aggregate` the types were regular pandas DataFrames, instead of the subclass.
Registering that `meta_nonempty` does keep it as `MyDataFrame` initially. I put some print statements in those methods to print the type of `inputs[0]` and `type(_concat(inputs))` and get
```
combine
aggregate
aggregate
aggregate
aggregate
split-out=2
```
So initially we're OK, but by the time we do the final `aggregate` we've lost the subclass.
Contributor guide
Research direction
Start by running the provided test.py reproduction and compare the computed types with and without split_out. Trace DecomposableGroupbyAggregation.combine and aggregate, focusing on where the final aggregate changes from MyDataFrame to pandas.DataFrame. Done means preserving the subclass for split_out=2 and adding a regression test covering both cases.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- pandas, python
- Domain
- data-engineering
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 42/100