[rtmalikian]

Fixes #9220.

Description

partition_dataset rounded each partition's size independently (int(r / rsum * data_len + 0.5)), so the rounded sizes could sum to less than len(data) and trailing items ended up in no partition — e.g. ratios [1, 1, 1] on 10 items gave [[0, 1, 2], [3, 4, 5], [6, 7, 8]] (item 9 lost), and [0.8, 0.1, 0.1] on 63 items gave sizes [50, 6, 6] (one case lost). The fix computes the rounded sizes first and, only when they fall short of data_len, assigns the remainder to the last partition, so every item lands in exactly one partition. All other behaviour is unchanged: the docstring example ([0.6, 0.2, 0.2] on 5 items) and the overshoot case (capped at data_len) behave exactly as before. partition_dataset_classes benefits automatically since it delegates to this function.

Types of changes

• Non-breaking change (fix or new feature that would not break existing functionality).
• Breaking change (fix or new feature that would cause existing functionality to change).
• New tests added to cover the changes.
• Integration tests passed locally by running ./runtests.sh -f -u --net --coverage.
• Quick tests passed locally by running the partition dataset unit tests (see Verification).
• In-line docstrings updated.
• Documentation updated, tested make html command in the docs/ folder.

Verification

• PYTHONPATH=. python -m pytest tests/data/test_partition_dataset.py tests/data/test_partition_dataset_classes.py -q: 16 passed (12 partition cases incl. 4 new, 4 classes cases).
• Negative control (fix stashed, new tests kept): exactly the 3 new loss-case tests fail; all 9 pre-existing partition tests pass.
• Issue repro post-fix: [0.8, 0.1, 0.1] on 63 items (shuffled, seed 0) gives sizes [50, 6, 7] with zero lost items; [1, 1, 1] on 10 items gives [[0, 1, 2], [3, 4, 5], [6, 7, 8, 9]].
• Exhaustive sweep post-fix: for ratios [0.8, 0.1, 0.1], [1, 1, 1], and [1, 1, 1, 1, 1], every length from 1 to 1000 partitions with no lost or duplicated items.
• ruff check clean on both touched files; ruff format --check clean.

Additional Notes

Author

Raphael Malikian, rtmalikian@gmail.com — If you'd like similar help with your own project or codebase, I'm open to collaborate — feel free to reach out.
This PR was prepared with AI assistance.

[coderabbitai]

Walkthrough

The ratio-based branch of partition_dataset now precomputes partition sizes and assigns any remainder to the last positive-ratio partition. The partition loop uses those sizes. New tests cover equal and fractional ratios across several dataset lengths, including a zero-ratio partition.

Priority: Normal

Estimated code review effort: 2 (Simple) | ~10 minutes

Severity of issue fixed: Medium

Merge Risk: Low · up to 5712e

The change stops items from being silently dropped when rounded ratio sizes fall short. In very small datasets the leftover items can all land in one partition, leaving others empty. This is uneven but not lossy, so the change is mergeable with awareness of the allocation tradeoff.

Pre-merge checks: 5

Check name	Status	Explanation
Title check		The title clearly identifies the main change: preventing item loss in partition_dataset caused by rounded ratio sizes.
Description check		The description explains the defect, implementation, compatibility behavior, tests, verification results, and change type. It is mostly complete; the full integration test was correctly marked as not …
Linked Issues check		Issue #9220 requires ratio-based partition_dataset to place every item in exactly one partition. The reviewed code computes rounded sizes before slicing and adds any shortfall to the last positive-r…
Out of Scope Changes check		The changes are limited to ratio-size handling in monai/data/utils.py and focused partition_dataset regression tests. These changes directly implement issue #9220. No unrelated change is shown.
Docstring Coverage		Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentat…

Full details: Description check

Explanation

The description explains the defect, implementation, compatibility behavior, tests, verification results, and change type. It is mostly complete; the full integration test was correctly marked as not run.

Full details: Linked Issues check

Explanation

Issue #9220 requires ratio-based partition_dataset to place every item in exactly one partition. The reviewed code computes rounded sizes before slicing and adds any shortfall to the last positive-ratio partition. Tests cover the reported 3-way loss cases, five ratios, the zero-ratio final partition, and the documented [0.6, 0.2, 0.2] example. partition_dataset_classes delegates to this function and receives the fix.

Full details: Docstring Coverage

Explanation

Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 1 functions across 2 files.

Finishing Touches

🧪 Generate unit tests (beta)

• Create a new PR

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

🔇 Additional comments: 1

monai/data/utils.py-1212-1212: 1

1212-1212: Minor · Functional Correctness · Quick win

⚠️ Unverified finding
Verification ran but could not confirm this finding. It is shown for review, not as a verified issue.

Keep zero-ratio partitions empty.

With 10 items and ratios [1, 1, 1, 0], sizes is [3, 3, 3, 0]. This line assigns the remaining item to the zero-ratio partition. Add the shortfall to the last partition with a positive ratio, and test this case. (raw.githubusercontent.com)

As per path instructions, “Examine code for logical error or inconsistencies, and suggest what may be changed to addressed these.”

Source: Path instructions

---

Review info

Run configuration

• Source mode: Auto
• Repository: Project-MONAI/MONAI
• Branch: dev
• File: .coderabbit.yaml
• Commit: e0cee14
• Layers: Repository YAML
• Review profile: CHILL
• Plan: Advanced
• Run ID: f8b38bb3-6ea5-44ca-8642-82e1304678f7

Commits

Reviewing files that changed from the base of the PR and between e0cee14 and afff1e3.

Files selected for processing: 2

• monai/data/utils.py
• tests/data/test_partition_dataset.py

Included review availability: This review used your included allowance. Your plan provides up to 8 included reviews per hour; 7 remain after this review.

[rtmalikian]

Good catch, thank you — you are right that the remainder could land in a zero-ratio partition. Fixed in commit 5712ee0: the rounding remainder now goes to the last partition with a positive ratio, so a zero-ratio partition stays empty. Added a regression test for ratios [1, 1, 1, 0] on 10 items; all 13 partition tests pass, and the new case fails on the pre-fix code.

[coderabbitai]

🔇 Additional comments: 1

monai/data/utils.py-1217-1217: 1

1217-1217: Major · Functional Correctness · Heavy lift

⚠️ Unverified finding
Verification ran but could not confirm this finding. It is shown for review, not as a verified issue.

Distribute the shortfall across partitions.

With two items and ratios [1, 1, 1, 1, 1, 0], this assigns both items to one partition, although two positive-ratio partitions could each receive one.

The effect also compounds in partition_dataset_classes: with one item per class and ratios [1, 1, 1, 0], every item goes to the third partition. The first two partitions are empty. Coordinate remainder allocation across classes as well as across partitions.

---

Review info

Run configuration

• Source mode: Auto
• Repository: Project-MONAI/MONAI
• Branch: dev
• File: .coderabbit.yaml
• Commit: e0cee14
• Layers: Repository YAML
• Review profile: CHILL
• Plan: Advanced
• Run ID: ede5ccc7-4008-433f-98ab-e8d52adfd6ac

Commits

Reviewing files that changed from the base of the PR and between afff1e3 and 5712ee0.

Files selected for processing: 2

• monai/data/utils.py
• tests/data/test_partition_dataset.py

Included review availability: This review used your included allowance. Your plan provides up to 8 included reviews per hour; 7 remain after this review.

[rtmalikian]

Good point, and thanks for checking the small-data cases — I verified the arithmetic: with ratios [1, 1, 1, 1, 1, 0] on 2 items, the current fix does put both items in the last positive-ratio partition, and the same rounding concentrates each single-item class in the third partition for partition_dataset_classes.

I kept this fix deliberately narrow — the reported issue was items being lost entirely, and the remainder-to-last-positive-partition rule guarantees every item lands in exactly one partition while keeping zero-ratio partitions empty. Spreading the shortfall (e.g. largest-remainder, or coordinating across classes) would change the allocation for the original cases too, so I'd rather not expand the scope on a bot suggestion alone.

Maintainers: would you like the remainder distributed across positive-ratio partitions within this PR, or is the no-loss guarantee enough here, with the balancing as a follow-up? Happy to implement whichever you prefer.