Skip to content

Commit

Permalink
Move num_groups computation out of if branches
Browse files Browse the repository at this point in the history
  • Loading branch information
adam-narozniak committed Aug 23, 2024
1 parent 7dfe5a6 commit db90c9d
Showing 1 changed file with 3 additions and 4 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -91,20 +91,20 @@ def _create_int_partition_id_to_natural_id(self) -> None:
unique_natural_ids = sorted(unique_natural_ids)
num_unique_natural_ids = len(unique_natural_ids)
remainder = num_unique_natural_ids % self._group_size
num_groups = num_unique_natural_ids // self._group_size
# Note that the number of groups might be different that this number
# due to certain modes, it's a base value.

if self._mode == "allow-bigger":
num_groups = num_unique_natural_ids // self._group_size
groups_of_natural_ids = np.array_split(unique_natural_ids, num_groups)
elif self._mode == "drop-reminder":
num_groups = num_unique_natural_ids // self._group_size
# Narrow down the unique_natural_ids to not have a bigger group
# which is the behavior of the np.array_split
unique_natural_ids = unique_natural_ids[
: int(num_groups * self._group_size)
]
groups_of_natural_ids = np.array_split(unique_natural_ids, num_groups)
elif self._mode == "allow-smaller":
num_groups = num_unique_natural_ids // self._group_size
if remainder > 0:
last_group_ids = unique_natural_ids[-remainder:]
unique_natural_ids = unique_natural_ids[
Expand All @@ -122,7 +122,6 @@ def _create_int_partition_id_to_natural_id(self) -> None:
f"enables strict mode or relax the mode parameter. Refer to the "
f"documentation of the mode parameter for the available modes."
)
num_groups = num_unique_natural_ids // self._group_size
groups_of_natural_ids = np.array_split(unique_natural_ids, num_groups)
else:
raise ValueError(
Expand Down

0 comments on commit db90c9d

Please sign in to comment.