Skip to content

Commit 59616c5

Browse files
authored
BUG: idxmin/max raising for arrow dtypes (#55368)
1 parent 06a54ad commit 59616c5

File tree

3 files changed

+28
-3
lines changed

3 files changed

+28
-3
lines changed

doc/source/whatsnew/v2.1.2.rst

+1
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ Fixed regressions
2222

2323
Bug fixes
2424
~~~~~~~~~
25+
- Fixed bug in :meth:`DataFrame.idxmin` and :meth:`DataFrame.idxmax` raising for arrow dtypes (:issue:`55368`)
2526
- Fixed bug in :meth:`DataFrame.resample` not respecting ``closed`` and ``label`` arguments for :class:`~pandas.tseries.offsets.BusinessDay` (:issue:`55282`)
2627
- Fixed bug in :meth:`DataFrame.resample` where bin edges were not correct for :class:`~pandas.tseries.offsets.BusinessDay` (:issue:`55281`)
2728
- Silence ``Period[B]`` warnings introduced by :issue:`53446` during normal plotting activity (:issue:`55138`)

pandas/core/arrays/arrow/array.py

+14-3
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,10 @@
3030
from pandas.util._decorators import doc
3131
from pandas.util._validators import validate_fillna_kwargs
3232

33-
from pandas.core.dtypes.cast import can_hold_element
33+
from pandas.core.dtypes.cast import (
34+
can_hold_element,
35+
infer_dtype_from_scalar,
36+
)
3437
from pandas.core.dtypes.common import (
3538
is_array_like,
3639
is_bool_dtype,
@@ -1624,13 +1627,21 @@ def _reduce(
16241627
pa_result = self._reduce_pyarrow(name, skipna=skipna, **kwargs)
16251628

16261629
if keepdims:
1627-
result = pa.array([pa_result.as_py()], type=pa_result.type)
1630+
if isinstance(pa_result, pa.Scalar):
1631+
result = pa.array([pa_result.as_py()], type=pa_result.type)
1632+
else:
1633+
result = pa.array(
1634+
[pa_result],
1635+
type=to_pyarrow_type(infer_dtype_from_scalar(pa_result)[0]),
1636+
)
16281637
return type(self)(result)
16291638

16301639
if pc.is_null(pa_result).as_py():
16311640
return self.dtype.na_value
1632-
else:
1641+
elif isinstance(pa_result, pa.Scalar):
16331642
return pa_result.as_py()
1643+
else:
1644+
return pa_result
16341645

16351646
def _explode(self):
16361647
"""

pandas/tests/frame/test_reductions.py

+13
Original file line numberDiff line numberDiff line change
@@ -1056,6 +1056,19 @@ def test_idxmax_numeric_only(self, numeric_only):
10561056
expected = Series([1, 0, 1], index=["a", "b", "c"])
10571057
tm.assert_series_equal(result, expected)
10581058

1059+
def test_idxmax_arrow_types(self):
1060+
# GH#55368
1061+
pytest.importorskip("pyarrow")
1062+
1063+
df = DataFrame({"a": [2, 3, 1], "b": [2, 1, 1]}, dtype="int64[pyarrow]")
1064+
result = df.idxmax()
1065+
expected = Series([1, 0], index=["a", "b"])
1066+
tm.assert_series_equal(result, expected)
1067+
1068+
result = df.idxmin()
1069+
expected = Series([2, 1], index=["a", "b"])
1070+
tm.assert_series_equal(result, expected)
1071+
10591072
def test_idxmax_axis_2(self, float_frame):
10601073
frame = float_frame
10611074
msg = "No axis named 2 for object type DataFrame"

0 commit comments

Comments
 (0)