Backport PR #55368 on branch 2.1.x (BUG: idxmin/max raising for arrow dtypes) (#55377)

BUG: idxmin/max raising for arrow dtypes (#55368)

(cherry picked from commit 59616c5)

Author: phofl

doc/source/whatsnew/v2.1.2.rst

@@ -21,7 +21,7 @@ Fixed regressions
 
 Bug fixes
 ~~~~~~~~~
--
+- Fixed bug in :meth:`DataFrame.idxmin` and :meth:`DataFrame.idxmax` raising for arrow dtypes (:issue:`55368`)
 -
 
 .. ---------------------------------------------------------------------------

pandas/core/arrays/arrow/array.py

@@ -30,6 +30,7 @@
 from pandas.util._decorators import doc
 from pandas.util._validators import validate_fillna_kwargs
 
+from pandas.core.dtypes.cast import infer_dtype_from_scalar
 from pandas.core.dtypes.common import (
     is_array_like,
     is_bool_dtype,
@@ -1595,13 +1596,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):
         """

pandas/tests/frame/test_reductions.py

@@ -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"