Skip to content

REF: implement Groupby idxmin, idxmax without fallback #38264

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Dec 4, 2020
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
44 changes: 43 additions & 1 deletion pandas/core/groupby/generic.py
Original file line number Diff line number Diff line change
Expand Up @@ -50,14 +50,14 @@
)
from pandas.core.dtypes.missing import isna, notna

from pandas.core import algorithms, nanops
from pandas.core.aggregation import (
agg_list_like,
aggregate,
maybe_mangle_lambdas,
reconstruct_func,
validate_func_kwargs,
)
import pandas.core.algorithms as algorithms
from pandas.core.arrays import Categorical, ExtensionArray
from pandas.core.base import DataError, SpecificationError
import pandas.core.common as com
Expand Down Expand Up @@ -1826,4 +1826,46 @@ def nunique(self, dropna: bool = True) -> DataFrame:
self._insert_inaxis_grouper_inplace(results)
return results

@Appender(DataFrame.idxmax.__doc__)
def idxmax(self, axis=0, skipna: bool = True):
axis = DataFrame._get_axis_number(axis)
numeric_only = None if axis == 0 else False

def func(df):
# NB: here we use numeric_only=None, in DataFrame it is False GH#38217
res = df._reduce(
nanops.nanargmax,
"argmax",
axis=axis,
skipna=skipna,
numeric_only=numeric_only,
)
indices = res._values
index = df._get_axis(axis)
result = [index[i] if i >= 0 else np.nan for i in indices]
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

isn't this a .take()?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it's take-like, but using take directly raises with Int64Index. This is almost identical to what we have in DataFrame.idxmax/idxmin.

return df._constructor_sliced(result, index=res.index)

return self._python_apply_general(func, self._obj_with_exclusions)

@Appender(DataFrame.idxmin.__doc__)
def idxmin(self, axis=0, skipna: bool = True):
axis = DataFrame._get_axis_number(axis)
numeric_only = None if axis == 0 else False

def func(df):
# NB: here we use numeric_only=None, in DataFrame it is False GH#38217
res = df._reduce(
nanops.nanargmin,
"argmin",
axis=axis,
skipna=skipna,
numeric_only=numeric_only,
)
indices = res._values
index = df._get_axis(axis)
result = [index[i] if i >= 0 else np.nan for i in indices]
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

same

return df._constructor_sliced(result, index=res.index)

return self._python_apply_general(func, self._obj_with_exclusions)

boxplot = boxplot_frame_groupby
19 changes: 1 addition & 18 deletions pandas/core/groupby/groupby.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,6 @@ class providing the base-class of operations.
import datetime
from functools import partial, wraps
import inspect
import re
import types
from typing import (
Callable,
Expand Down Expand Up @@ -797,23 +796,7 @@ def curried(x):
if name in base.plotting_methods:
return self.apply(curried)

try:
return self._python_apply_general(curried, self._obj_with_exclusions)
except TypeError as err:
if not re.search(
"reduction operation '.*' not allowed for this dtype", str(err)
):
# We don't have a cython implementation
# TODO: is the above comment accurate?
raise

if self.obj.ndim == 1:
# this can be called recursively, so need to raise ValueError
raise ValueError

# GH#3688 try to operate item-by-item
result = self._aggregate_item_by_item(name, *args, **kwargs)
return result
return self._python_apply_general(curried, self._obj_with_exclusions)

wrapper.__name__ = name
return wrapper
Expand Down
21 changes: 21 additions & 0 deletions pandas/tests/groupby/test_function.py
Original file line number Diff line number Diff line change
Expand Up @@ -521,6 +521,27 @@ def test_idxmin_idxmax_returns_int_types(func, values):
tm.assert_frame_equal(result, expected)


def test_idxmin_idxmax_axis1():
df = DataFrame(np.random.randn(10, 4), columns=["A", "B", "C", "D"])
df["A"] = [1, 2, 3, 1, 2, 3, 1, 2, 3, 4]

gb = df.groupby("A")

res = gb.idxmax(axis=1)

alt = df.iloc[:, 1:].idxmax(axis=1)
indexer = res.index.get_level_values(1)

tm.assert_series_equal(alt[indexer], res.droplevel("A"))

df["E"] = pd.date_range("2016-01-01", periods=10)
gb2 = df.groupby("A")

msg = "reduction operation 'argmax' not allowed for this dtype"
with pytest.raises(TypeError, match=msg):
gb2.idxmax(axis=1)


def test_groupby_cumprod():
# GH 4095
df = DataFrame({"key": ["b"] * 10, "value": 2})
Expand Down