Skip to content

Commit

Permalink
FIX-#2596: revert groupby changes
Browse files Browse the repository at this point in the history
Signed-off-by: Dmitry Chigarev <dmitry.chigarev@intel.com>
  • Loading branch information
dchigarev committed Feb 9, 2021
1 parent 74ebfab commit 0cd090a
Showing 1 changed file with 14 additions and 9 deletions.
23 changes: 14 additions & 9 deletions modin/pandas/groupby.py
Original file line number Diff line number Diff line change
Expand Up @@ -622,9 +622,18 @@ def corr(self, **kwargs):
return self._default_to_pandas(lambda df: df.corr(**kwargs))

def fillna(self, **kwargs):
result = self._apply_agg_function(
lambda df: df.fillna(**kwargs), overwrite_groupby_kwargs={"as_index": True}
new_groupby_kwargs = self._kwargs.copy()
new_groupby_kwargs["as_index"] = True
work_object = type(self)(
df=self._df,
by=self._by,
axis=self._axis,
idx_name=self._idx_name,
drop=self._drop,
squeeze=self._squeeze,
**new_groupby_kwargs,
)
result = work_object._apply_agg_function(lambda df: df.fillna(**kwargs))
# pandas does not name the index on fillna
result._query_compiler.set_index_name(None)
return result
Expand Down Expand Up @@ -894,7 +903,7 @@ def _wrap_aggregation(
return result.squeeze()
return result

def _apply_agg_function(self, f, overwrite_groupby_kwargs=None, *args, **kwargs):
def _apply_agg_function(self, f, *args, **kwargs):
"""
Perform aggregation and combine stages based on a given function.
Expand All @@ -904,8 +913,6 @@ def _apply_agg_function(self, f, overwrite_groupby_kwargs=None, *args, **kwargs)
----------
f: callable
The function to apply to each group.
overwrite_groupby_kwargs: dict (optional),
GroupBy kwargs to overwrite.
Returns
-------
Expand All @@ -914,17 +921,15 @@ def _apply_agg_function(self, f, overwrite_groupby_kwargs=None, *args, **kwargs)
assert callable(f) or isinstance(
f, dict
), "'{0}' object is not callable and not a dict".format(type(f))
groupby_kwargs = self._kwargs.copy()
if overwrite_groupby_kwargs is not None:
groupby_kwargs.update(overwrite_groupby_kwargs)

new_manager = self._query_compiler.groupby_agg(
by=self._by,
is_multi_by=self._is_multi_by,
axis=self._axis,
agg_func=f,
agg_args=args,
agg_kwargs=kwargs,
groupby_kwargs=groupby_kwargs,
groupby_kwargs=self._kwargs,
drop=self._drop,
)
if self._idx_name is not None and self._as_index:
Expand Down

0 comments on commit 0cd090a

Please sign in to comment.