Skip to content

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

Open
jzhan-2026 wants to merge 10 commits into
apache:masterfrom
jzhan-2026:pandas-rank-params
Open

[SPARK-58971][PS] Add pct, na_option, and axis parameters to Series.rank and DataFrame.rank#58254
jzhan-2026 wants to merge 10 commits into
apache:masterfrom
jzhan-2026:pandas-rank-params

Conversation

@jzhan-2026

@jzhan-2026 jzhan-2026 commented Aug 24, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

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

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.

>>> 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).

@jzhan-2026

jzhan-2026 commented Aug 24, 2026

Copy link
Copy Markdown
Contributor Author

@holdenk Could you help take a look on the PR when you have time? Thanks.

@jzhan-2026 jzhan-2026 changed the title SPARK-58971][PS] Add pct, na_option, and axis parameters to Series.rank and DataFrame.rank [SPARK-58971][PS] Add pct, na_option, and axis parameters to Series.rank and DataFrame.rank Aug 24, 2026
@holdenk

holdenk commented Aug 24, 2026

Copy link
Copy Markdown
Contributor

I'll take a look, there are some failing "slow" pandas tests, they look like they could be bad tests but if you can also take a look @jzhan-2026 that would be great!

holdenk
holdenk reviewed Aug 24, 2026
View reviewed changes
Comment thread python/pyspark/pandas/series.py Outdated
Comment on lines +4363 to +4383
# 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)

@holdenk holdenk Aug 24, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think the new ordering is a little screwed up, I took a little of time but not 100% sure. Tentatively my guess is making the asc_func and then applying it to the natural order column name would work better. I think it could also simplify the code a bit since the multiple different calls to .desc() / .desc_nulls_last / desc_nulls... etc.

@jzhan-2026 jzhan-2026 Aug 25, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You are right about this; the way that I added the na_option ordering is confusing. I have refactored the code so it looks cleaner and less error-prone.

But I noticed another thing - there was a preexisting ordering bug in the code, I added a test case in the code and verified the bug is real. I filed a JIRA to track it here: https://issues.apache.org/jira/browse/SPARK-59011

Let me know if you think we should fix this in the same PR or prioritize the bug fix soon in a new PR.

@holdenk holdenk Aug 25, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah I think we can defer that to a follow up provided we do that ~soon.

@holdenk

holdenk commented Aug 24, 2026

Copy link
Copy Markdown
Contributor

Did a bit more of a poke, I think there might be something off in the new ordering code not just the ttests.

@jzhan-2026
jzhan-2026 requested a review from holdenk August 25, 2026 22:43
holdenk
holdenk reviewed Aug 26, 2026
View reviewed changes

@holdenk holdenk left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Overall LGTM, I do want @devin-petersohn to double check the how we're handling the new params situation here.

Comment thread python/pyspark/pandas/frame.py
Comment on lines +11764 to +11765
na_option: Literal["keep", "top", "bottom"] = "keep",
pct: bool = False,

@holdenk holdenk Aug 26, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@devin-petersohn is this how we're handling adding the missing params normally? Or do we shove them after **kwargs?

@devin-petersohn devin-petersohn Aug 27, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In the past we've put new parameters at the end.

Another option is to require them as kwargs with the * notation.

@devin-petersohn devin-petersohn Aug 27, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The reason we put them at the end was in case users are using positional notation for their args. This won't be fully backwards compatible if we put the new args in the middle.

@jzhan-2026 jzhan-2026 Aug 27, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for your comments! I have moved the newly added args at the end to ensure backward compatibility.

Or do we shove them after **kwargs?

I personally lean towards putting args at the end because using **kwargs diverges from all other recently-added params in this codebase (SPARK-46163, SPARK-47997, SPARK-53645 — none used *); also creates an asymmetry with pandas' own signature where these params aren't keyword-only. Happy to go a different direction if you have a preference though!

holdenk
holdenk approved these changes Aug 27, 2026
View reviewed changes

@holdenk holdenk left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@holdenk holdenk holdenk approved these changes
@devin-petersohn devin-petersohn Awaiting requested review from devin-petersohn

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@jzhan-2026 @holdenk @devin-petersohn