Skip to content

Commit

Permalink
BUG: idxmin/max raising for arrow dtypes (#55368)
Browse files Browse the repository at this point in the history
  • Loading branch information
phofl authored Oct 3, 2023
1 parent 06a54ad commit 59616c5
Show file tree
Hide file tree
Showing 3 changed files with 28 additions and 3 deletions.
1 change: 1 addition & 0 deletions doc/source/whatsnew/v2.1.2.rst
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ Fixed regressions

Bug fixes
~~~~~~~~~
- Fixed bug in :meth:`DataFrame.idxmin` and :meth:`DataFrame.idxmax` raising for arrow dtypes (:issue:`55368`)
- Fixed bug in :meth:`DataFrame.resample` not respecting ``closed`` and ``label`` arguments for :class:`~pandas.tseries.offsets.BusinessDay` (:issue:`55282`)
- Fixed bug in :meth:`DataFrame.resample` where bin edges were not correct for :class:`~pandas.tseries.offsets.BusinessDay` (:issue:`55281`)
- Silence ``Period[B]`` warnings introduced by :issue:`53446` during normal plotting activity (:issue:`55138`)
Expand Down
17 changes: 14 additions & 3 deletions pandas/core/arrays/arrow/array.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,10 @@
from pandas.util._decorators import doc
from pandas.util._validators import validate_fillna_kwargs

from pandas.core.dtypes.cast import can_hold_element
from pandas.core.dtypes.cast import (
can_hold_element,
infer_dtype_from_scalar,
)
from pandas.core.dtypes.common import (
is_array_like,
is_bool_dtype,
Expand Down Expand Up @@ -1624,13 +1627,21 @@ def _reduce(
pa_result = self._reduce_pyarrow(name, skipna=skipna, **kwargs)

if keepdims:
result = pa.array([pa_result.as_py()], type=pa_result.type)
if isinstance(pa_result, pa.Scalar):
result = pa.array([pa_result.as_py()], type=pa_result.type)
else:
result = pa.array(
[pa_result],
type=to_pyarrow_type(infer_dtype_from_scalar(pa_result)[0]),
)
return type(self)(result)

if pc.is_null(pa_result).as_py():
return self.dtype.na_value
else:
elif isinstance(pa_result, pa.Scalar):
return pa_result.as_py()
else:
return pa_result

def _explode(self):
"""
Expand Down
13 changes: 13 additions & 0 deletions pandas/tests/frame/test_reductions.py
Original file line number Diff line number Diff line change
Expand Up @@ -1056,6 +1056,19 @@ def test_idxmax_numeric_only(self, numeric_only):
expected = Series([1, 0, 1], index=["a", "b", "c"])
tm.assert_series_equal(result, expected)

def test_idxmax_arrow_types(self):
# GH#55368
pytest.importorskip("pyarrow")

df = DataFrame({"a": [2, 3, 1], "b": [2, 1, 1]}, dtype="int64[pyarrow]")
result = df.idxmax()
expected = Series([1, 0], index=["a", "b"])
tm.assert_series_equal(result, expected)

result = df.idxmin()
expected = Series([2, 1], index=["a", "b"])
tm.assert_series_equal(result, expected)

def test_idxmax_axis_2(self, float_frame):
frame = float_frame
msg = "No axis named 2 for object type DataFrame"
Expand Down

0 comments on commit 59616c5

Please sign in to comment.