BUG: Preserve Series/DataFrame subclasses through groupby operations

[JBGreisman]

This pull request fixes a bug that caused subclassed Series/DataFrames to revert to pd.Series and pd.DataFrame after groupby() operations. This is a follow-up on a previously abandoned pull request (#28573)

• closes groupby and resample methods do not preserve subclassed data structures #28330
• tests added / passed
• passes black pandas
• passes git diff upstream/master -u -- "*.py" | flake8 --diff
• whatsnew entry

[JBGreisman]

I updated my branch with additional tests, and a revised entry in the whatsnew rst file.

Please let me know if there are any additional recommendations for how to proceed with this PR.

[JBGreisman]

Thanks for the mypy advice -- I managed to refactor that section a bit to only call cast() if needed and avoid calling the constructor again.

The one additional change made here is to make sure GroupBy.size(), GroupBy.ngroup(), and GroupBy.cumcount() preserve subclassed types. While you had previously mentioned just keeping these return types as Series, I think these changes minimize the possibility for surprise when subclassing DataFrame/Series by ensuring that groupby reductions/transformations preserve the subclasses.

[JBGreisman]

@jreback -- please let me know if there are any additional changes to this PR that you think are necessary.

[jreback]

small comment, ping on green.

[JBGreisman]

@jreback I made the recommend changes:

• I reverted the changes in DataFrameGroupBy._cython_agg_blocks() because they did not prove necessary to pass my tests.
• I added a GroupBy._constructor property to simplify the calls in ngroup, cumcount, and size.

[simonjayhawkins]

@JBGreisman generally lgtm. one question.

[JBGreisman]

@simonjayhawkins Thanks for the comment -- I agree that line should have used self.obj._constructor, and changed it

[jreback]

hmm this should be FrameOrSeries I think.

why is there an else here?

[JBGreisman]

I had thought this should be a Series because if self.obj is a DataFrame, it returns self.obj._constructor_sliced. Do you think I should name this property _series_constructor or something comparable to make that behavior more apparent?

The else was there because mypy was complaining about a missing return statement. I can restructure this with an assertion to avoid an else statement and keep mypy from complaining:

@property
def _series_constructor(self) -> Type["Series"]:
    # GH28330 preserve subclassed Series/DataFrames                                                                                                                                                   
    if isinstance(self.obj, DataFrame):
        return self.obj._constructor_sliced
    assert isinstance(self.obj, Series)
    return self.obj._constructor

Please let me know if you have a better way to structure this.

[jreback]

yes pls use an assertion

maybe @simonjayhawkins or @WillAyd can help with the annotation itself

[JBGreisman]

thanks -- I changed _constructor to use the assertion as above. I think it would also make sense to change the name of the property to _series_constructor in order to clarify the return type, but I'll hold off for additional comments.

[WillAyd]

I think the annotation is correct given the way it is implemented, though I am not sure about the implementation. Why do we need to dispatch to constructor_sliced for DataFrames? Seems slightly unnatural to have to force that

[JBGreisman]

This ends up being used in GroupBy.ngroup(), GroupBy.cumcount(), and GroupBy.size(). Prior to these changes, these methods had returned Series, regardless of whether they were called from a SeriesGroupBy or DataFrameGroupBy object. I had updated this to still return a Series-type while preserving the subclasses -- to avoid things reverting back to pd.Series if they were called from a subclassed DataFrame/Series.

As such, the motivation for dispatching to constructor_sliced for DataFrames was to avoid changing the return type for these different GroupBy methods from their prior behavior. Do you think it would make sense to make a larger change here that alters these functions to have different return types if called from SeriesGroupBy vs. DataFrameGroupBy?

[WillAyd]

Hmm that's unfortunate... Can you rename this property to _obj_constructor instead? I find the current name a little confusing

[WillAyd]

Or maybe even _obj_1d_constructor to be even more explicit. This is definitely for special cases so want to signal as such

[JBGreisman]

yeah -- I think _obj_1d_constructor is most clear. I'll make the changes.

[JBGreisman]

As per the discussion above, the GroupBy property has been renamed to _obj_1d_constructor to clarify its expected return type, and I refactored it to avoid using an else statement.

[jreback]

thanks @JBGreisman very nice, thanks for sticking with it and being responsive!

[JBGreisman]

no worries -- thanks for the help/suggestions!