using _mgr apply with 2 failing tests

Author: sharkipelago

pandas/core/algorithms.py

@@ -1088,7 +1088,6 @@ def rank(
         )
     else:
         raise TypeError("Array with ndim > 2 are not supported.")
-
     return ranks
 
 

pandas/core/arrays/base.py

@@ -2408,6 +2408,7 @@ def _rank(
         """
         See Series.rank.__doc__.
         """
+
         if axis != 0:
             raise NotImplementedError
 

pandas/core/generic.py

@@ -9275,34 +9275,25 @@ def rank(
             msg = "na_option must be one of 'keep', 'top', or 'bottom'"
             raise ValueError(msg)
 
-        def ranker(data):
-            if data.ndim == 2:
-                # i.e. DataFrame, we cast to ndarray
-                values = data.values
-            else:
-                # i.e. Series, can dispatch to EA
-                values = data._values
-
-            if isinstance(values, ExtensionArray):
-                ranks = values._rank(
-                    axis=axis_int,
+        def ranker(blk_values):
+            if isinstance(blk_values, ExtensionArray) and blk_values.ndim == 1:
+                ranks = blk_values._rank(
+                    axis=0,
                     method=method,
                     ascending=ascending,
                     na_option=na_option,
                     pct=pct,
                 )
             else:
                 ranks = algos.rank(
-                    values,
-                    axis=axis_int,
+                    blk_values,
+                    axis=1 - axis_int,
                     method=method,
                     ascending=ascending,
                     na_option=na_option,
                     pct=pct,
                 )
-
-            ranks_obj = self._constructor(ranks, **data._construct_axes_dict())
-            return ranks_obj.__finalize__(self, method="rank")
+            return ranks
 
         if numeric_only:
             if self.ndim == 1 and not is_numeric_dtype(self.dtype):
@@ -9315,7 +9306,10 @@ def ranker(data):
         else:
             data = self
 
-        return ranker(data)
+        result = data._mgr.apply(ranker)
+        return self._constructor_from_mgr(result, axes=result.axes).__finalize__(
+            self, method="rank"
+        )
 
     @doc(_shared_docs["compare"], klass=_shared_doc_kwargs["klass"])
     def compare(

pandas/tests/frame/methods/test_rank.py

@@ -405,6 +405,7 @@ def test_rank_inf_and_nan(self, contents, dtype, frame_or_series):
         # Shuffle the testing array and expected results in the same way
         random_order = np.random.default_rng(2).permutation(len(values))
         obj = frame_or_series(values[random_order])
+        print("TYPE", type(obj))
         expected = frame_or_series(exp_order[random_order], dtype="float64")
         result = obj.rank()
         tm.assert_equal(result, expected)
@@ -498,3 +499,38 @@ def test_rank_string_dtype(self, string_dtype_no_object):
             exp_dtype = "float64"
         expected = Series([1, 2, None, 3], dtype=exp_dtype)
         tm.assert_series_equal(result, expected)
+
+    @pytest.mark.parametrize(
+        "method,og_dtype,expected_dtype",
+        [
+            ("average", "UInt32", "Float64"),
+            ("average", "Float32", "Float64"),
+            ("average", "int32[pyarrow]", "double[pyarrow]"),
+            ("min", "Int32", "Float64"),
+            ("min", "Float32", "Float64"),
+            ("min", "int32[pyarrow]", "double[pyarrow]"),
+        ],
+    )
+    def test_rank_extension_array_dtype(self, method, og_dtype, expected_dtype):
+        # GH#52829
+        result = DataFrame([4, 89, 33], dtype=og_dtype).rank()
+        if method == "average":
+            expected = DataFrame([1.0, 3.0, 2.0], dtype=expected_dtype)
+        else:
+            expected = DataFrame([1, 3, 2], dtype=expected_dtype)
+        tm.assert_frame_equal(result, expected)
+
+    def test_rank_mixed_extension_array_dtype(self):
+        result = DataFrame(
+            {
+                "base": Series([4, 5, 6]),
+                "extension": Series([7, 8, 9], dtype="int32[pyarrow]"),
+            }
+        ).rank(method="min")
+        expected = DataFrame(
+            {
+                "base": Series([1.0, 2.0, 3.0], dtype="float64"),
+                "extension": Series([1, 2, 3], dtype="uint64[pyarrow]"),
+            }
+        )
+        tm.assert_frame_equal(result, expected)