2d extension arrays

Author: sharkipelago

pandas/core/arrays/base.py

@@ -2415,9 +2415,6 @@ def _rank(
         See Series.rank.__doc__.
         """
 
-        if axis != 0:
-            raise NotImplementedError
-
         return rank(
             self,
             axis=axis,

pandas/core/generic.py

@@ -9277,33 +9277,27 @@ def rank(
             raise ValueError(msg)
 
         def ranker(blk_values):
-            if isinstance(blk_values, ExtensionArray) and blk_values.ndim == 1:
+            if axis_int == 0:
+                blk_values = blk_values.T
+            if isinstance(blk_values, ExtensionArray):
                 ranks = blk_values._rank(
-                    axis=0,
+                    axis=axis_int,
                     method=method,
                     ascending=ascending,
                     na_option=na_option,
                     pct=pct,
                 )
             else:
-                if axis_int == 0:
-                    ranks = algos.rank(
-                        blk_values.T,
-                        axis=axis_int,
-                        method=method,
-                        ascending=ascending,
-                        na_option=na_option,
-                        pct=pct,
-                    ).T
-                else:
-                    ranks = algos.rank(
-                        blk_values,
-                        axis=axis_int,
-                        method=method,
-                        ascending=ascending,
-                        na_option=na_option,
-                        pct=pct,
-                    )
+                ranks = algos.rank(
+                    blk_values,
+                    axis=axis_int,
+                    method=method,
+                    ascending=ascending,
+                    na_option=na_option,
+                    pct=pct,
+                )
+            if axis_int == 0:
+                ranks = ranks.T
             return ranks
 
         if numeric_only:

pandas/tests/frame/methods/test_rank.py

@@ -16,6 +16,8 @@
     DataFrame,
     Index,
     Series,
+    to_datetime,
+    to_timedelta,
 )
 import pandas._testing as tm
 
@@ -511,37 +513,84 @@ def test_rank_string_dtype(self, string_dtype_no_object):
                 "double[pyarrow]",
                 marks=td.skip_if_no("pyarrow"),
             ),
-            ("min", "Int32", "Float64"),
-            ("min", "Float32", "Float64"),
+            ("min", "Int32", "UInt64"),
+            ("min", "Float32", "UInt64"),
             pytest.param(
                 "min",
                 "int32[pyarrow]",
-                "double[pyarrow]",
+                "uint64[pyarrow]",
                 marks=td.skip_if_no("pyarrow"),
             ),
         ],
     )
     def test_rank_extension_array_dtype(self, method, og_dtype, expected_dtype):
         # GH#52829
-        result = DataFrame([4, 89, 33], dtype=og_dtype).rank()
+        result = DataFrame([4, 89, 33], dtype=og_dtype).rank(method=method)
         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):
+        # GH#52829
         pytest.importorskip("pyarrow")
         result = DataFrame(
             {
                 "base": Series([4, 5, 6]),
-                "extension": Series([7, 8, 9], dtype="int32[pyarrow]"),
+                "pyarrow": 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]"),
+                "pyarrow": Series([1, 2, 3], dtype="uint64[pyarrow]"),
             }
         )
         tm.assert_frame_equal(result, expected)
+
+    def test_2d_extension_array_datetime(self):
+        # GH#52829
+        df = DataFrame(
+            {
+                "year": to_datetime(["2012-1-1", "2013-1-1", "2014-1-1"]),
+                "week": to_datetime(["2012-1-2", "2012-1-9", "2012-1-16"]),
+                "day": to_datetime(["2012-1-3", "2012-1-4", "2012-1-5"]),
+            }
+        )
+        axis0_expected = DataFrame(
+            {"year": [1.0, 2.0, 3.0], "week": [1.0, 2.0, 3.0], "day": [1.0, 2.0, 3.0]}
+        )
+        axis1_expected = DataFrame(
+            {"year": [1.0, 3.0, 3.0], "week": [2.0, 2.0, 2.0], "day": [3.0, 1.0, 1.0]}
+        )
+        tm.assert_frame_equal(df.rank(), axis0_expected)
+        tm.assert_frame_equal(df.rank(1), axis1_expected)
+
+    def test_2d_extension_array_timedelta(self):
+        # GH#52829
+        df = DataFrame(
+            {
+                "day": to_timedelta(["0 days", "1 day", "2 days"]),
+                "hourly": to_timedelta(["23 hours", "24 hours", "25 hours"]),
+                "minute": to_timedelta(
+                    ["1439 minutes", "1440 minutes", "1441 minutes"]
+                ),
+            }
+        )
+        axis0_expected = DataFrame(
+            {
+                "day": [1.0, 2.0, 3.0],
+                "hourly": [1.0, 2.0, 3.0],
+                "minute": [1.0, 2.0, 3.0],
+            }
+        )
+        axis1_expected = DataFrame(
+            {
+                "day": [1.0, 2.0, 3.0],
+                "hourly": [2.0, 2.0, 2.0],
+                "minute": [3.0, 2.0, 1.0],
+            }
+        )
+        tm.assert_frame_equal(df.rank(), axis0_expected)
+        tm.assert_frame_equal(df.rank(1), axis1_expected)