API: rank with nullable dtypes preserve NA

Author: jbrockmendel

doc/source/whatsnew/v3.0.0.rst

@@ -95,6 +95,7 @@ Other enhancements
 - Support passing a :class:`Iterable[Hashable]` input to :meth:`DataFrame.drop_duplicates` (:issue:`59237`)
 - Support reading Stata 102-format (Stata 1) dta files (:issue:`58978`)
 - Support reading Stata 110-format (Stata 7) dta files (:issue:`47176`)
+- :meth:`Series.rank` and :meth:`DataFrame.rank` with numpy-nullable dtypes preserve ``NA`` values and return ``UInt64`` dtype where appropriate instead of casting ``NA`` to ``NaN`` with ``float64`` dtype (:issue:`??`)
 
 .. ---------------------------------------------------------------------------
 .. _whatsnew_300.notable_bug_fixes:

pandas/core/arrays/masked.py

@@ -12,6 +12,7 @@
 import numpy as np
 
 from pandas._libs import (
+    algos as libalgos,
     lib,
     missing as libmissing,
 )
@@ -992,6 +993,49 @@ def copy(self) -> Self:
         mask = self._mask.copy()
         return self._simple_new(data, mask)
 
+    def _rank(
+        self,
+        *,
+        axis: AxisInt = 0,
+        method: str = "average",
+        na_option: str = "keep",
+        ascending: bool = True,
+        pct: bool = False,
+    ):
+        # Avoid going through copy-making ensure_data in algorithms.rank
+        if axis != 0 or self.ndim != 1:
+            raise NotImplementedError
+
+        from pandas.core.arrays import FloatingArray
+
+        data = self._data
+        if data.dtype.kind == "b":
+            data = data.view("uint8")
+
+        result = libalgos.rank_1d(
+            data,
+            is_datetimelike=False,
+            ties_method=method,
+            ascending=ascending,
+            na_option=na_option,
+            pct=pct,
+            mask=self.isna(),
+        )
+        if na_option in ["top", "bottom"]:
+            mask = np.zeros(self.shape, dtype=bool)
+        else:
+            mask = self._mask.copy()
+
+        if method != "average" and not pct:
+            if na_option not in ["top", "bottom"]:
+                result[self._mask] = 0  # avoid warning on casting
+            result = result.astype("uint64", copy=False)
+            from pandas.core.arrays import IntegerArray
+
+            return IntegerArray(result, mask=mask)
+
+        return FloatingArray(result, mask=mask)
+
     @doc(ExtensionArray.duplicated)
     def duplicated(
         self, keep: Literal["first", "last", False] = "first"

pandas/tests/series/methods/test_rank.py

@@ -68,6 +68,11 @@ def expected_dtype(dtype, method, pct=False):
             exp_dtype = "double[pyarrow]"
         else:
             exp_dtype = "uint64[pyarrow]"
+    elif dtype in ["Float64", "Int64"]:
+        if method == "average" or pct:
+            exp_dtype = "Float64"
+        else:
+            exp_dtype = "UInt64"
 
     return exp_dtype
 
@@ -257,7 +262,7 @@ def test_rank_nullable_integer(self):
         exp = Series([None, 2, None, 3, 3, 2, 3, 1], dtype="Int64")
         result = exp.rank(na_option="keep")
 
-        expected = Series([np.nan, 2.5, np.nan, 5.0, 5.0, 2.5, 5.0, 1.0])
+        expected = Series([None, 2.5, None, 5.0, 5.0, 2.5, 5.0, 1.0], dtype="Float64")
 
         tm.assert_series_equal(result, expected)
 
@@ -302,6 +307,11 @@ def test_rank_tie_methods_on_infs_nans(
                 exp_dtype = "float64[pyarrow]"
             else:
                 exp_dtype = "uint64[pyarrow]"
+        elif dtype == "Float64":
+            if rank_method == "average":
+                exp_dtype = "Float64"
+            else:
+                exp_dtype = "UInt64"
         else:
             exp_dtype = "float64"
 
@@ -327,7 +337,8 @@ def test_rank_tie_methods_on_infs_nans(
         result = iseries.rank(
             method=rank_method, na_option=na_option, ascending=ascending
         )
-        tm.assert_series_equal(result, Series(expected, dtype=exp_dtype))
+        exp_ser = Series(expected, dtype=exp_dtype)
+        tm.assert_series_equal(result, exp_ser)
 
     def test_rank_desc_mix_nans_infs(self):
         # GH 19538
@@ -439,7 +450,7 @@ def test_rank_ea_small_values(self):
             dtype="Float64",
         )
         result = ser.rank(method="min")
-        expected = Series([4, 1, 3, np.nan, 2])
+        expected = Series([4, 1, 3, NA, 2], dtype="UInt64")
         tm.assert_series_equal(result, expected)