From afff1e3a04b8eb92b1ab1a00c2ec5260166e9911 Mon Sep 17 00:00:00 2001 From: Raphael Malikian Date: Tue, 27 Oct 2026 16:32:07 -0700 Subject: [PATCH 1/2] fix(data): keep every item in exactly one partition when ratios are rounded partition_dataset rounded each partition size independently, so the rounded sizes could sum to less than len(data) and trailing items ended up in no partition. Assign any rounding remainder to the last partition. Fixes #9220 Signed-off-by: Raphael Malikian --- monai/data/utils.py | 11 ++++- tests/data/test_partition_dataset.py | 69 +++++++++++++++++++++++++++- 2 files changed, 76 insertions(+), 4 deletions(-) diff --git a/monai/data/utils.py b/monai/data/utils.py index d548ed72486..ea60f421244 100644 --- a/monai/data/utils.py +++ b/monai/data/utils.py @@ -1203,9 +1203,16 @@ def partition_dataset( if ratios: next_idx = 0 rsum = sum(ratios) - for r in ratios: + sizes = [int(r / rsum * data_len + 0.5) for r in ratios] + if sum(sizes) < data_len: + # Independent rounding can leave trailing items in no partition + # (e.g. ratios [1, 1, 1] on 10 items rounds to 3 + 3 + 3); assign + # the remainder to the last partition so every item lands in + # exactly one partition. + sizes[-1] += data_len - sum(sizes) + for size in sizes: start_idx = next_idx - next_idx = min(start_idx + int(r / rsum * data_len + 0.5), data_len) + next_idx = min(start_idx + size, data_len) datasets.append([data[i] for i in indices[start_idx:next_idx]]) return datasets diff --git a/tests/data/test_partition_dataset.py b/tests/data/test_partition_dataset.py index c93a6c76821..a8cd7c285b2 100644 --- a/tests/data/test_partition_dataset.py +++ b/tests/data/test_partition_dataset.py @@ -117,10 +117,75 @@ ] -class TestPartitionDataset(unittest.TestCase): +TEST_CASE_9 = [ + { + "data": list(range(10)), + "ratios": [1, 1, 1], + "num_partitions": None, + "shuffle": False, + "seed": 0, + "drop_last": True, + "even_divisible": True, + }, + [[0, 1, 2], [3, 4, 5], [6, 7, 8, 9]], +] + +TEST_CASE_10 = [ + { + "data": list(range(63)), + "ratios": [0.8, 0.1, 0.1], + "num_partitions": None, + "shuffle": False, + "seed": 0, + "drop_last": True, + "even_divisible": True, + }, + [list(range(50)), list(range(50, 56)), list(range(56, 63))], +] +TEST_CASE_11 = [ + { + "data": list(range(32)), + "ratios": [1, 1, 1, 1, 1], + "num_partitions": None, + "shuffle": False, + "seed": 0, + "drop_last": True, + "even_divisible": True, + }, + [list(range(6)), list(range(6, 12)), list(range(12, 18)), list(range(18, 24)), list(range(24, 32))], +] + +TEST_CASE_12 = [ + { + "data": [1, 2, 3, 4, 5], + "ratios": [0.6, 0.2, 0.2], + "num_partitions": None, + "shuffle": False, + "seed": 0, + "drop_last": True, + "even_divisible": True, + }, + [[1, 2, 3], [4], [5]], +] + + +class TestPartitionDataset(unittest.TestCase): @parameterized.expand( - [TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4, TEST_CASE_5, TEST_CASE_6, TEST_CASE_7, TEST_CASE_8] + [ + TEST_CASE_1, + TEST_CASE_2, + TEST_CASE_3, + TEST_CASE_4, + TEST_CASE_5, + TEST_CASE_6, + TEST_CASE_7, + TEST_CASE_8, + TEST_CASE_9, + TEST_CASE_10, + TEST_CASE_11, + TEST_CASE_12, + ] ) def test_value(self, input_param, result): self.assertListEqual(partition_dataset(**input_param), result) From 5712ee01beeb45cf5c29f52e95d40f0f17a0696b Mon Sep 17 00:00:00 2001 From: Raphael Malikian Date: Wed, 28 Oct 2026 09:24:41 -0700 Subject: [PATCH 2/2] fix(data): keep zero-ratio partitions empty when assigning the rounding remainder CodeRabbit review on #9230 noted the remainder assignment goes to the last partition even when its ratio is 0 (ratios [1, 1, 1, 0] on 10 items rounds to 3 + 3 + 3 + 0, and the leftover item landed in the empty zero-ratio partition). The remainder now goes to the last partition with a positive ratio instead. Regression test added (TEST_CASE_13); 13/13 partition tests pass, and the new case fails on the pre-fix code. Signed-off-by: Raphael Malikian --- monai/data/utils.py | 12 +++++++++--- tests/data/test_partition_dataset.py | 14 ++++++++++++++ 2 files changed, 23 insertions(+), 3 deletions(-) diff --git a/monai/data/utils.py b/monai/data/utils.py index ea60f421244..d85b241d6da 100644 --- a/monai/data/utils.py +++ b/monai/data/utils.py @@ -1207,9 +1207,15 @@ def partition_dataset( if sum(sizes) < data_len: # Independent rounding can leave trailing items in no partition # (e.g. ratios [1, 1, 1] on 10 items rounds to 3 + 3 + 3); assign - # the remainder to the last partition so every item lands in - # exactly one partition. - sizes[-1] += data_len - sum(sizes) + # the remainder so every item lands in exactly one partition. + # It goes to the last partition with a positive ratio: a + # zero-ratio partition must stay empty (e.g. ratios + # [1, 1, 1, 0] on 10 items rounds to 3 + 3 + 3 + 0, and the + # leftover item belongs in the third partition, not the empty one). + for i in range(len(ratios) - 1, -1, -1): + if ratios[i] > 0: + sizes[i] += data_len - sum(sizes) + break for size in sizes: start_idx = next_idx next_idx = min(start_idx + size, data_len) diff --git a/tests/data/test_partition_dataset.py b/tests/data/test_partition_dataset.py index a8cd7c285b2..57e34e71b7b 100644 --- a/tests/data/test_partition_dataset.py +++ b/tests/data/test_partition_dataset.py @@ -169,6 +169,19 @@ [[1, 2, 3], [4], [5]], ] +TEST_CASE_13 = [ + { + "data": list(range(10)), + "ratios": [1, 1, 1, 0], + "num_partitions": None, + "shuffle": False, + "seed": 0, + "drop_last": True, + "even_divisible": True, + }, + [[0, 1, 2], [3, 4, 5], [6, 7, 8, 9], []], +] + class TestPartitionDataset(unittest.TestCase): @parameterized.expand( @@ -185,6 +198,7 @@ class TestPartitionDataset(unittest.TestCase): TEST_CASE_10, TEST_CASE_11, TEST_CASE_12, + TEST_CASE_13, ] ) def test_value(self, input_param, result):