Generalize NormalizedPartitioning class#22246
Merged
rapids-bot[bot] merged 5 commits intorapidsai:mainfrom Apr 23, 2026
Merged
Conversation
rjzamora
added
2 - In Progress
github-actions
Bot
added
Python
rjzamora
commented
Apr 21, 2026
Comment on lines
-654
to
+662
| inter_rank_modulus: int | ||
| """The inter-rank modulus.""" | ||
| inter_rank_indices: tuple[int, ...] | ||
| """The inter-rank column indices.""" | ||
| local_modulus: int | None | ||
| """The local modulus.""" | ||
| local_indices: tuple[int, ...] | ||
| """The local column indices.""" | ||
| inter_rank_scheme: HashScheme | None | ||
| local_scheme: HashScheme | None | Literal["inherit"] |
Member
Author
There was a problem hiding this comment.
This is the central change here. Everything else is mostly needed to align with this change/simplification.
rjzamora
commented
Apr 21, 2026
pentschev
approved these changes
Apr 22, 2026
Member
pentschev
left a comment
pentschev
left a comment
There was a problem hiding this comment.
I think this looks good, left only some nits. Thanks Rick.
Comment on lines
+691
to
+699
| lhs_loc = self.local_scheme | ||
| rhs_loc = other.local_scheme | ||
| if type(lhs_loc) is not type(rhs_loc): | ||
| return False | ||
| if isinstance(lhs_loc, HashScheme) and isinstance(rhs_loc, HashScheme): | ||
| return lhs_loc.modulus == rhs_loc.modulus and len( | ||
| lhs_loc.column_indices | ||
| ) == len(rhs_loc.column_indices) | ||
| return True |
Member
There was a problem hiding this comment.
Suggested change
| lhs_loc = self.local_scheme | |
| rhs_loc = other.local_scheme | |
| if type(lhs_loc) is not type(rhs_loc): | |
| return False | |
| if isinstance(lhs_loc, HashScheme) and isinstance(rhs_loc, HashScheme): | |
| return lhs_loc.modulus == rhs_loc.modulus and len( | |
| lhs_loc.column_indices | |
| ) == len(rhs_loc.column_indices) | |
| return True | |
| if lhs_loc != rhs_loc: | |
| if not (isinstance(lhs_loc, HashScheme) and isinstance(rhs_loc, HashScheme)): | |
| return False | |
| return lhs_loc.modulus == rhs_loc.modulus and len(lhs_loc.column_indices) == len(rhs_loc.column_indices) | |
| return True |
Member
There was a problem hiding this comment.
I think this is generally easier to read and doesn't rely on identity of type() in favor of a proper isinstance().
Member
Author
There was a problem hiding this comment.
I tried to resolve the logic, but did it a bit differently.
rjzamora
force-pushed
the
normalized-partitioning-rewrite
branch
from
5c40703 to
1865865
Compare
rjzamora
added
3 - Ready for Review
rjzamora
added
5 - Ready to Merge
Member
Author
|
/merge |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
The
NormalizedPartitioningis veryHashSchemespecific. It promotes usage that focuses on the effective hash modulus used for inter-rank and local partitioning. When we adoptOrderScheme(rapidsai/rapidsmpf#853), we will need to compare other descriptions (e.g. sort boundaries). This PR revises theNormalizedPartitioningto reflect the reality of the underlyingPartitioning"scheme" more directly.Checklist