Pandas: groupby and resample methods do not preserve subclassed data structures

[grge]

Code sample

class MyDataFrame(pd.DataFrame):
    @property
    def _constructor(self):
        return MyDataFrame

dates = pd.date_range('2019', freq='H', periods=1000)
my_df = MyDataFrame(np.arange(len(dates)), index=dates)

print(type(my_df)) 
# __main__.MyDataFrame (✓)

print(type(my_df.diff())) 
# __main__.MyDataFrame (✓)

print(type(my_df.sample(1))) 
# __main__.MyDataFrame (✓)

print(type(my_df.rolling('5H').mean()))
# __main__.MyDataFrame (✓)

print(type(my_df.groupby(my_df.index.dayofweek).mean()))
# pandas.core.frame.DataFrame (✘)

print(type(my_df.resample('D').mean()))
# pandas.core.frame.DataFrame (✘)

Problem description

Originally posted on SO.

The intended behaviour for chain-able methods on subclassed data structures is clearly that the operation returns an instance of the subclass (i.e. MyDataFrame), rather than the native type (i.e. DataFrame). This is the current behaviour for most operations (e.g., slicing, sampling, sorting) but not resample and groupby.

Currently groupby and resample both return explicitly constructed pandas datatypes, e.g. here:

https://github.com/pandas-dev/pandas/blob/ac693331400bce4747c63aa76c53c6e3488933ae/pandas/core/groupby/generic.py#L338

To get the expected behaviour, the intermediary classes (e.g. DataFrameGroupBy) would need to retain information about the calling class so that the appropriate constructor can be used (i.e. one of _constructor or _constructor_sliced or _constructor_expanddim).

Note that operations that use Window and Rolling already appear have the expected behaviour because these assemble their results via a call to concat such as this one:

https://github.com/pandas-dev/pandas/blob/171c71611886aab8549a8620c5b0071a129ad685/pandas/core/window.py#L325

Output of pd.show_versions()

INSTALLED VERSIONS

commit : None
python : 3.7.4.final.0
python-bits : 64
OS : Windows
OS-release : 10
machine : AMD64
processor : Intel64 Family 6 Model 142 Stepping 9, GenuineIntel
byteorder : little
LC_ALL : None
LANG : None
LOCALE : None.None

pandas : 0.25.1
numpy : 1.17.1
pytz : 2019.2
dateutil : 2.8.0
pip : 19.2.3
setuptools : 40.8.0
Cython : None
pytest : None
hypothesis : None
sphinx : None
blosc : None
feather : None
xlsxwriter : None
lxml.etree : None
html5lib : None
pymysql : None
psycopg2 : None
jinja2 : 2.10.1
IPython : 7.8.0
pandas_datareader: None
bs4 : None
bottleneck : None
fastparquet : None
gcsfs : None
lxml.etree : None
matplotlib : 3.1.1
numexpr : 2.7.0
odfpy : None
openpyxl : None
pandas_gbq : None
pyarrow : None
pytables : None
s3fs : None
scipy : 1.3.1
sqlalchemy : None
tables : 3.5.2
xarray : None
xlrd : None
xlwt : None
xlsxwriter : None

[alkasm]

Worth noting that the group-by objects do store the object it started with in the attribute obj (which makes sense, since it has to aggregate the data from it), so those constructors could indeed be called like they are elsewhere in the library. I just implemented this using those methods; only requires a change to groupby/generic.py! Here's it in use:

import pandas as pd

class MySeries(pd.Series):
    pass

class MyDataFrame(pd.DataFrame):
    @property
    def _constructor(self):
        return MyDataFrame
    _constructor_sliced = MySeries

MySeries._constructor_expanddim = MyDataFrame

for cls in (pd.DataFrame, MyDataFrame):
    df = cls(
        {"a": reversed(range(10)), "b": list('aaaabbbccc')}
    )
    s = df.groupby("b").sum()
    print(type(df))
    print(type(s))
    print(type(s['a']))

<class 'pandas.core.frame.DataFrame'>
<class 'pandas.core.frame.DataFrame'>
<class 'pandas.core.series.Series'>
<class '__main__.MyDataFrame'>
<class '__main__.MyDataFrame'>
<class '__main__.MySeries'>

Try it out from my fork here: https://github.com/alkasm/pandas/tree/groupby-preserve-subclass

I haven't taken a look at the resampling source code at all. However, it seems to use groupby (not surprisingly, since I'd bet it uses a Grouper to do its thing), at least in the example you gave, as it rightfully returns the subclass now:

import pandas as pd
import numpy as np

class MyDataFrame(pd.DataFrame):
    @property
    def _constructor(self):
        return MyDataFrame

dates = pd.date_range('2019', freq='H', periods=1000)
my_df = MyDataFrame(np.arange(len(dates)), index=dates)

print(type(my_df))
# <class '__main__.MyDataFrame'> (✓)

print(type(my_df.diff()))
# <class '__main__.MyDataFrame'> (✓)

print(type(my_df.sample(1)))
# <class '__main__.MyDataFrame'> (✓)

print(type(my_df.rolling('5H').mean()))
# <class '__main__.MyDataFrame'> (✓)

print(type(my_df.groupby(my_df.index.dayofweek).mean()))
# <class '__main__.MyDataFrame'> (✓)

print(type(my_df.resample('D').mean()))
# <class '__main__.MyDataFrame'> (✓)

I will open up a PR after I'm able to look into the resampling stuff a little more and confirm whether or not this covers the bases.

[grge]

That was fast! I'll try out the fork when I've got access to a less locked down PC.

As far as I can tell this solves the problem perfectly.

[alkasm]

AFAICT, resampling will do the right thing, as it just applies the functions/classes from the groupby module, so I don't think anything special is necessary. There isn't really hardcoded Series or DataFrame there so I think we're good. I will submit the PR shortly. Thanks for the issue @grge.

Edit: PR submitted: https://github.com/pandas-dev/pandas/pull/28573