[SPARK-58971][PS] Add pct, na_option, and axis parameters to Series.rank and DataFrame.rank

python/pyspark/pandas/frame.py

@@ -11756,12 +11756,13 @@ def _result_aggregated(
         # dtype: bool
         return first_series(DataFrame(internal))
 
-    # TODO(SPARK-46167): add pct, na_option parameter
     def rank(
         self,
         method: Literal["average", "min", "max", "first", "dense"] = "average",
         ascending: bool = True,
         numeric_only: bool = False,
+        na_option: str = "keep",
+        pct: bool = False,
         axis: Axis = 0,
     ) -> "DataFrame":
         """
@@ -11789,6 +11790,18 @@ def rank(
             .. versionchanged:: 4.0.0
                 The default value of ``numeric_only`` is now ``False``.
 
+        na_option : {'keep', 'top', 'bottom'}, default 'keep'
+            * keep: leave NA values where they are
+            * top: smallest rank if ascending
+            * bottom: largest rank if ascending
+
+            .. versionadded:: 4.4.0
+
+        pct : bool, default False
+            Whether or not to display the returned rankings in percentile form.
+
+            .. versionadded:: 4.4.0
+
         axis : {0 or 'index', 1 or 'columns'}, default 0
             Axis along which to rank:
 
@@ -11798,7 +11811,6 @@ def rank(
             .. note:: For axis=1, pandas UDF is used which may have performance overhead
                 for very wide DataFrames (100+ columns).
 
-
         Returns
         -------
         ranks : same type as caller
@@ -11887,14 +11899,24 @@ def rank(
 
         if axis == 0:
             return psdf._apply_series_op(
-                lambda psser: psser._rank(method=method, ascending=ascending), should_resolve=True
+                lambda psser: psser._rank(
+                    method=method, ascending=ascending, na_option=na_option, pct=pct
+                ),
+                should_resolve=True,
             )
         else:
             # Fast path for small dataframes
             limit = get_option("compute.shortcut_limit")
             pdf = psdf.head(limit + 1)._to_internal_pandas()
             if len(pdf) <= limit:
-                pdf_rank = pdf.rank(method=method, ascending=ascending, axis=1, numeric_only=False)
+                pdf_rank = pdf.rank(
+                    method=method,
+                    ascending=ascending,
+                    axis=1,
+                    numeric_only=False,
+                    na_option=na_option,
+                    pct=pct,
+                )
                 return DataFrame(InternalFrame.from_pandas(pdf_rank))
 
             column_label_strings = [
@@ -11911,9 +11933,13 @@ def rank(
             )
             def rank_axis_1(*cols: pd.Series) -> pd.DataFrame:
                 pdf_row = pd.concat(cols, axis=1, keys=column_label_strings)
-                return pdf_row.rank(method=method, ascending=ascending, axis=1).rename(
-                    columns=dict(zip(pdf_row.columns, column_label_strings))
-                )
+                return pdf_row.rank(
+                    method=method,
+                    ascending=ascending,
+                    axis=1,
+                    na_option=na_option,
+                    pct=pct,
+                ).rename(columns=dict(zip(pdf_row.columns, column_label_strings)))
 
             ranked_struct_col = rank_axis_1(*psdf._internal.data_spark_columns)
             new_data_columns = [

python/pyspark/pandas/series.py

@@ -4190,9 +4190,14 @@ def quantile(psser: Series) -> PySparkColumn:
 
             return self._reduce_for_stat_function(quantile, name="quantile")
 
-    # TODO: add axis, pct, na_option parameter
     def rank(
-        self, method: str = "average", ascending: bool = True, numeric_only: bool = False
+        self,
+        method: str = "average",
+        ascending: bool = True,
+        numeric_only: bool = False,
+        na_option: str = "keep",
+        pct: bool = False,
+        axis: int = 0,
     ) -> "Series":
         """
         Compute numerical data ranks (1 through n) along axis. Equal values are
@@ -4219,6 +4224,20 @@ def rank(
             .. versionchanged:: 4.0.0
                 The default value of ``numeric_only`` is now ``False``.
 
+        na_option : {'keep', 'top', 'bottom'}, default 'keep'
+            * keep: leave NA values where they are
+            * top: smallest rank if ascending
+            * bottom: largest rank if ascending
+
+            .. versionadded:: 4.4.0
+
+        pct : bool, default False
+            Whether or not to display the returned rankings in percentile form.
+
+            .. versionadded:: 4.4.0
+
+        axis : {0 or 'index'}, default 0
+            Parameter needed for compatibility with DataFrame.
 
         Returns
         -------
@@ -4286,45 +4305,89 @@ def rank(
         y    b
         z    c
         Name: A, dtype: object
+
+        With pct=True, ranks are expressed as percentiles.
+
+        >>> s = ps.Series([1, 2, 2, 3], name='A')
+        >>> s.rank(pct=True)
+        0    0.25
+        1    0.625
+        2    0.625
+        3    1.0
+        Name: A, dtype: float64
+
+        With na_option='top', NaN values are assigned the smallest rank.
+
+        >>> s = ps.Series([1, float('nan'), 2, 3], name='A')
+        >>> s.rank(na_option='top')
+        0    2.0
+        1    1.0
+        2    3.0
+        3    4.0
+        Name: A, dtype: float64
+
+        With na_option='bottom', NaN values are assigned the largest rank.
+
+        >>> s.rank(na_option='bottom')
+        0    1.0
+        1    4.0
+        2    2.0
+        3    3.0
+        Name: A, dtype: float64
         """
+        validate_axis(axis)
         is_numeric = isinstance(self.spark.data_type, (NumericType, BooleanType))
         if numeric_only and not is_numeric:
             raise TypeError("Series.rank does not allow numeric_only=True with non-numeric dtype.")
         else:
-            return self._rank(method, ascending).spark.analyzed
+            return self._rank(method, ascending, na_option=na_option, pct=pct).spark.analyzed
 
     def _rank(
         self,
         method: str = "average",
         ascending: bool = True,
         *,
         part_cols: Sequence["ColumnOrName"] = (),
+        na_option: str = "keep",
+        pct: bool = False,
     ) -> "Series":
         if method not in ["average", "min", "max", "first", "dense"]:
             msg = "method must be one of 'average', 'min', 'max', 'first', 'dense'"
             raise ValueError(msg)
+        if na_option not in ["keep", "top", "bottom"]:
+            raise ValueError("na_option must be one of 'keep', 'top', 'bottom'")
 
         if self._internal.index_level > 1:
             raise NotImplementedError("rank do not support MultiIndex now")
 
+        # Determine ordering with null placement based on na_option.
+        # 'top' always assigns the smallest rank to NaN, 'bottom' the largest,
+        # regardless of ascending direction.
         if ascending:
-            asc_func = PySparkColumn.asc
+            sort_col = (
+                self.spark.column.asc_nulls_first()
+                if na_option == "top"
+                else self.spark.column.asc_nulls_last()
+            )
+            nat_order_col = F.col(NATURAL_ORDER_COLUMN_NAME).asc()
         else:
-            asc_func = PySparkColumn.desc
+            sort_col = (
+                self.spark.column.desc_nulls_first()
+                if na_option == "top"
+                else self.spark.column.desc_nulls_last()
+            )
+            nat_order_col = F.col(NATURAL_ORDER_COLUMN_NAME).desc()
 
         if method == "first":
             window = (
-                Window.orderBy(
-                    asc_func(self.spark.column),
-                    asc_func(F.col(NATURAL_ORDER_COLUMN_NAME)),
-                )
+                Window.orderBy(sort_col, nat_order_col)
                 .partitionBy(*part_cols)
                 .rowsBetween(Window.unboundedPreceding, Window.currentRow)
             )
             scol = F.row_number().over(window)
         elif method == "dense":
             window = (
-                Window.orderBy(asc_func(self.spark.column))
+                Window.orderBy(sort_col)
                 .partitionBy(*part_cols)
                 .rowsBetween(Window.unboundedPreceding, Window.currentRow)
             )
@@ -4337,7 +4400,7 @@ def _rank(
             elif method == "max":
                 stat_func = F.max
             window1 = (
-                Window.orderBy(asc_func(self.spark.column))
+                Window.orderBy(sort_col)
                 .partitionBy(*part_cols)
                 .rowsBetween(Window.unboundedPreceding, Window.currentRow)
             )
@@ -4346,6 +4409,26 @@ def _rank(
                 cast("List[ColumnOrName]", [self.spark.column]) + list(part_cols)
             ).rowsBetween(Window.unboundedPreceding, Window.unboundedFollowing)
             scol = stat_func(F.row_number().over(window1)).over(window2)
+
+        if pct:
+            partition_window = Window.partitionBy(*part_cols)
+            if method == "dense":
+                # For dense ranking, pct denominator is the number of distinct values.
+                if na_option == "keep":
+                    denom = F.max(F.when(self.spark.column.isNotNull(), scol)).over(
+                        partition_window
+                    )
+                else:
+                    denom = F.max(scol).over(partition_window)
+            elif na_option == "keep":
+                denom = F.count(self.spark.column).over(partition_window)
+            else:
+                denom = F.count(F.lit(1)).over(partition_window)
+            scol = scol / denom.cast(DoubleType())
+
+        if na_option == "keep":
+            scol = F.when(self.spark.column.isNotNull(), scol)
+
         return self._with_new_scol(scol.cast(DoubleType()))
 
     def filter(

python/pyspark/pandas/tests/computation/test_compute.py

@@ -360,6 +360,60 @@ def test_rank_axis(self):
         with self.assertRaisesRegex(ValueError, "No axis named"):
             psdf.rank(axis=2)
 
+    def test_rank_pct_na_option(self):
+        pdf = pd.DataFrame(
+            {"A": [1, 2, np.nan, 3], "B": [4, np.nan, 2, 1]},
+            columns=["A", "B"],
+        )
+        psdf = ps.from_pandas(pdf)
+
+        # pct=True
+        self.assert_eq(pdf.rank(pct=True).sort_index(), psdf.rank(pct=True).sort_index())
+
+        # na_option='top'
+        self.assert_eq(
+            pdf.rank(na_option="top").sort_index(), psdf.rank(na_option="top").sort_index()
+        )
+
+        # na_option='bottom'
+        self.assert_eq(
+            pdf.rank(na_option="bottom").sort_index(), psdf.rank(na_option="bottom").sort_index()
+        )
+
+        # pct + na_option combined
+        self.assert_eq(
+            pdf.rank(pct=True, na_option="top").sort_index(),
+            psdf.rank(pct=True, na_option="top").sort_index(),
+        )
+        self.assert_eq(
+            pdf.rank(pct=True, na_option="bottom").sort_index(),
+            psdf.rank(pct=True, na_option="bottom").sort_index(),
+        )
+
+        # all methods
+        for method in ["average", "min", "max", "first", "dense"]:
+            self.assert_eq(
+                pdf.rank(method=method, pct=True).sort_index(),
+                psdf.rank(method=method, pct=True).sort_index(),
+            )
+            self.assert_eq(
+                pdf.rank(method=method, na_option="bottom").sort_index(),
+                psdf.rank(method=method, na_option="bottom").sort_index(),
+            )
+
+        # axis=1 with pct and na_option
+        self.assert_eq(
+            pdf.rank(axis=1, pct=True).sort_index(), psdf.rank(axis=1, pct=True).sort_index()
+        )
+        self.assert_eq(
+            pdf.rank(axis=1, na_option="top").sort_index(),
+            psdf.rank(axis=1, na_option="top").sort_index(),
+        )
+
+        # invalid na_option
+        with self.assertRaisesRegex(ValueError, "na_option must be one of"):
+            psdf.rank(na_option="invalid")
+
     def test_nunique(self):
         pdf = pd.DataFrame({"A": [1, 2, 3], "B": [np.nan, 3, np.nan]}, index=np.random.rand(3))
         psdf = ps.from_pandas(pdf)

python/pyspark/pandas/tests/series/test_stat.py

@@ -344,6 +344,30 @@ def test_rank(self):
         with self.assertRaisesRegex(ValueError, msg):
             psser.rank(method="nothing")
 
+        # pct=True
+        pser = pd.Series([1, 2, 2, 3], name="x")
+        psser = ps.from_pandas(pser)
+        self.assert_eq(pser.rank(pct=True), psser.rank(pct=True).sort_index())
+
+        # na_option
+        pser = pd.Series([1, float("nan"), 2, 3], name="x")
+        psser = ps.from_pandas(pser)
+        self.assert_eq(pser.rank(na_option="top"), psser.rank(na_option="top").sort_index())
+        self.assert_eq(pser.rank(na_option="bottom"), psser.rank(na_option="bottom").sort_index())
+        self.assert_eq(
+            pser.rank(pct=True, na_option="top"),
+            psser.rank(pct=True, na_option="top").sort_index(),
+        )
+
+        # axis=0 is accepted (no-op for Series)
+        pser = pd.Series([1, 2, 3, 1], name="x")
+        psser = ps.from_pandas(pser)
+        self.assert_eq(pser.rank(axis=0), psser.rank(axis=0).sort_index())
+
+        # invalid na_option
+        with self.assertRaisesRegex(ValueError, "na_option must be one of"):
+            psser.rank(na_option="bad")
+
         msg = "method must be one of 'average', 'min', 'max', 'first', 'dense'"
         with self.assertRaisesRegex(ValueError, msg):
             psser.rank(method="nothing")