This is an automated email from the ASF dual-hosted git repository.
asf-gitbox-commits pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new c696cba6eb2a [SPARK-58971][PS] Add pct, na_option, and axis parameters
to Series.rank and DataFrame.rank
c696cba6eb2a is described below
commit c696cba6eb2ae543c44df02208385a3cc5867e29
Author: jzhan-2026 <[email protected]>
AuthorDate: Wed Oct 28 17:21:20 2026 -0700
[SPARK-58971][PS] Add pct, na_option, and axis parameters to Series.rank
and DataFrame.rank
### What changes were proposed in this pull request?
Add `na_option`, `pct`, and `axis` parameters to `Series.rank()`, and
`na_option` and `pct` to `DataFrame.rank()`, to match the pandas API.
- `na_option` (`'keep'`/`'top'`/`'bottom'`): controls how `NaN` values are
ranked
- `pct`: expresses ranks as percentile fractions in `(0, 1]`
- `axis` on `Series.rank()`: accepts `0`/`'index'` only, added for API
compatibility
Co-authored with jing.zhansnowflake.com
### Why are the changes needed?
The pandas-on-Spark implementations were missing these parameters, making
it harder to migrate pandas code to Spark without modification.
`DataFrame.rank` was tracked in SPARK-46167; this PR also covers `Series.rank`.
### Does this PR introduce _any_ user-facing change?
Yes. Two new keyword parameters on `Series.rank()` and `DataFrame.rank()`.
Defaults are unchanged so existing code is unaffected.
```python
>>> ps.Series([1, float('nan'), 2]).rank(na_option='top')
0 2.0
1 1.0
2 3.0
>>> ps.Series([1, 2, 2, 3]).rank(pct=True)
0 0.25
1 0.625
2 0.625
3 1.0
```
### How was this patch tested?
New test method test_rank_pct_na_option in test_compute.py and new cases in
test_stat.py covering all methods, both axes, combined parameters, and invalid
input errors.
### Was this patch authored or co-authored using generative AI tooling?
This patch was co-authored with AI assistance (Claude Sonnet 4.6).
Closes #58254 from jzhan-2026/pandas-rank-params.
Authored-by: jzhan-2026 <[email protected]>
Signed-off-by: Holden Karau <[email protected]>
---
python/pyspark/pandas/frame.py | 38 ++++++--
python/pyspark/pandas/series.py | 101 ++++++++++++++++++---
.../pandas/tests/computation/test_compute.py | 64 +++++++++++++
python/pyspark/pandas/tests/series/test_stat.py | 32 +++++++
4 files changed, 218 insertions(+), 17 deletions(-)
diff --git a/python/pyspark/pandas/frame.py b/python/pyspark/pandas/frame.py
index 2dbb6f5f46b0..baebb079bd3e 100644
--- a/python/pyspark/pandas/frame.py
+++ b/python/pyspark/pandas/frame.py
@@ -11762,13 +11762,14 @@ defaultdict(<class 'list'>, {'col..., 'col...})]
# 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,
axis: Axis = 0,
+ na_option: Literal["keep", "top", "bottom"] = "keep",
+ pct: bool = False,
) -> "DataFrame":
"""
Compute numerical data ranks (1 through n) along axis. Equal values are
@@ -11804,6 +11805,17 @@ defaultdict(<class 'list'>, {'col..., 'col...})]
.. note:: For axis=1, pandas UDF is used which may have
performance overhead
for very wide DataFrames (100+ columns).
+ 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
Returns
-------
@@ -11893,14 +11905,24 @@ defaultdict(<class 'list'>, {'col..., 'col...})]
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 = [
@@ -11917,9 +11939,13 @@ defaultdict(<class 'list'>, {'col..., 'col...})]
)
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 = [
diff --git a/python/pyspark/pandas/series.py b/python/pyspark/pandas/series.py
index 7ce31bb6c092..0f32e467f08d 100644
--- a/python/pyspark/pandas/series.py
+++ b/python/pyspark/pandas/series.py
@@ -4195,9 +4195,14 @@ class Series(Frame, IndexOpsMixin, Generic[T]):
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,
+ axis: int = 0,
+ na_option: Literal["keep", "top", "bottom"] = "keep",
+ pct: bool = False,
) -> "Series":
"""
Compute numerical data ranks (1 through n) along axis. Equal values are
@@ -4224,6 +4229,20 @@ class Series(Frame, IndexOpsMixin, Generic[T]):
.. versionchanged:: 4.0.0
The default value of ``numeric_only`` is now ``False``.
+ axis : {0 or 'index'}, default 0
+ Parameter needed for compatibility with DataFrame.
+
+ 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
Returns
-------
@@ -4291,12 +4310,42 @@ class Series(Frame, IndexOpsMixin, Generic[T]):
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.250
+ 1 0.625
+ 2 0.625
+ 3 1.000
+ 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').sort_index()
+ 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').sort_index()
+ 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,
@@ -4304,32 +4353,42 @@ class Series(Frame, IndexOpsMixin, Generic[T]):
ascending: bool = True,
*,
part_cols: Sequence["ColumnOrName"] = (),
+ na_option: Literal["keep", "top", "bottom"] = "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")
+ # 'top' always assigns rank 1 to NaN (nulls first in window order).
+ # 'bottom' always assigns the largest rank to NaN (nulls last). 'keep'
also nulls last.
+ nulls_first = na_option == "top"
if ascending:
- asc_func = PySparkColumn.asc
+ null_func = (
+ PySparkColumn.asc_nulls_first if nulls_first else
PySparkColumn.asc_nulls_last
+ )
else:
- asc_func = PySparkColumn.desc
+ null_func = (
+ PySparkColumn.desc_nulls_first if nulls_first else
PySparkColumn.desc_nulls_last
+ )
+ sort_col = null_func(self.spark.column)
+ nat_order_col = F.col(NATURAL_ORDER_COLUMN_NAME).asc()
if method == "first":
window = (
- Window.orderBy(
- asc_func(self.spark.column),
- F.col(NATURAL_ORDER_COLUMN_NAME).asc(),
- )
+ 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)
)
@@ -4342,7 +4401,7 @@ class Series(Frame, IndexOpsMixin, Generic[T]):
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)
)
@@ -4351,6 +4410,26 @@ class Series(Frame, IndexOpsMixin, Generic[T]):
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(
diff --git a/python/pyspark/pandas/tests/computation/test_compute.py
b/python/pyspark/pandas/tests/computation/test_compute.py
index 394d46173142..c3b991bb55b2 100644
--- a/python/pyspark/pandas/tests/computation/test_compute.py
+++ b/python/pyspark/pandas/tests/computation/test_compute.py
@@ -363,6 +363,70 @@ class FrameComputeMixin:
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(),
+ )
+
+ # na_option with ascending=False
+ self.assert_eq(
+ pdf.rank(na_option="top", ascending=False).sort_index(),
+ psdf.rank(na_option="top", ascending=False).sort_index(),
+ )
+ self.assert_eq(
+ pdf.rank(na_option="bottom", ascending=False).sort_index(),
+ psdf.rank(na_option="bottom", ascending=False).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)
diff --git a/python/pyspark/pandas/tests/series/test_stat.py
b/python/pyspark/pandas/tests/series/test_stat.py
index 10cf9ed9e3ef..bfa09c899a2b 100644
--- a/python/pyspark/pandas/tests/series/test_stat.py
+++ b/python/pyspark/pandas/tests/series/test_stat.py
@@ -348,6 +348,38 @@ class SeriesStatMixin:
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(),
+ )
+ self.assert_eq(
+ pser.rank(na_option="top", ascending=False),
+ psser.rank(na_option="top", ascending=False).sort_index(),
+ )
+ self.assert_eq(
+ pser.rank(na_option="bottom", ascending=False),
+ psser.rank(na_option="bottom", ascending=False).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")
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
