Skip to content

Commit afd8491

Browse files
committed
Call __finalize__ in groupby iterator
This solution does not call `__finalize__` immediately on object creation in the `DataSplitter` instance. This is to allow `.apply` to avoid the overhead of calling `__finalize__` for intermediate objects.
1 parent 8f7673e commit afd8491

File tree

1 file changed

+12
-3
lines changed

1 file changed

+12
-3
lines changed

pandas/core/groupby/ops.py

+12-3
Original file line numberDiff line numberDiff line change
@@ -140,9 +140,16 @@ def get_iterator(
140140
splitter = self._get_splitter(data, axis=axis)
141141
keys = self._get_group_keys()
142142
for key, (i, group) in zip(keys, splitter):
143-
yield key, group
143+
yield key, group.__finalize__(data, method="groupby")
144144

145145
def _get_splitter(self, data: FrameOrSeries, axis: int = 0) -> "DataSplitter":
146+
"""
147+
Returns
148+
-------
149+
Generator yielding subsetted objects
150+
151+
__finalize__ has not been called for the the subsetted objects returned.
152+
"""
146153
comp_ids, _, ngroups = self.group_info
147154
return get_splitter(data, comp_ids, ngroups, axis=axis)
148155

@@ -921,7 +928,8 @@ class SeriesSplitter(DataSplitter):
921928
def _chop(self, sdata: Series, slice_obj: slice) -> Series:
922929
# fastpath equivalent to `sdata.iloc[slice_obj]`
923930
mgr = sdata._mgr.get_slice(slice_obj)
924-
return type(sdata)(mgr, name=sdata.name, fastpath=True)
931+
# __finalize__ not called here, must be applied by caller if applicable
932+
return sdata._constructor(mgr, name=sdata.name, fastpath=True)
925933

926934

927935
class FrameSplitter(DataSplitter):
@@ -937,7 +945,8 @@ def _chop(self, sdata: DataFrame, slice_obj: slice) -> DataFrame:
937945
# else:
938946
# return sdata.iloc[:, slice_obj]
939947
mgr = sdata._mgr.get_slice(slice_obj, axis=1 - self.axis)
940-
return type(sdata)(mgr)
948+
# __finalize__ not called here, must be applied by caller if applicable
949+
return sdata._constructor(mgr)
941950

942951

943952
def get_splitter(

0 commit comments

Comments
 (0)