Fix partition_dataset losing items when rounded ratio sizes fall short

monai/data/utils.py

@@ -1203,9 +1203,22 @@ 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 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 + 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
 

tests/data/test_partition_dataset.py

@@ -117,10 +117,89 @@
 ]
 
 
-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]],
+]
+
+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(
-        [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,
+            TEST_CASE_13,
+        ]
     )
     def test_value(self, input_param, result):
         self.assertListEqual(partition_dataset(**input_param), result)