dask / dask/dask-expr

DataFrame subclass lost in `groupby.agg` with `split_out` set.

Open
#1,024 1 comment 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.