fix

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>

Author: 0oshowero0

tests/test_metadata.py

@@ -215,11 +215,11 @@ def test_batch_meta_chunk_by_partition(self):
                 name="test_field", dtype=torch.float32, shape=(2,), production_status=ProductionStatus.READY_FOR_CONSUME
             )
         }
-        samples = [SampleMeta(partition_id=f"partition_{i % 4}", global_index=i, fields=fields) for i in range(10)]
+        samples = [SampleMeta(partition_id=f"partition_{i % 4}", global_index=i + 10, fields=fields) for i in range(10)]
         batch = BatchMeta(
             samples=samples,
-            custom_meta={i: {"uid": i} for i in range(10)},
-            _custom_backend_meta={i: {"test_field": {"dtype": torch.float32}} for i in range(10)},
+            custom_meta={i + 10: {"uid": i + 10} for i in range(10)},
+            _custom_backend_meta={i + 10: {"test_field": {"dtype": torch.float32}} for i in range(10)},
         )
 
         # Chunk according to partition_id
@@ -228,30 +228,30 @@ def test_batch_meta_chunk_by_partition(self):
         assert len(chunks) == 4
         assert len(chunks[0]) == 3
         assert chunks[0].partition_ids == ["partition_0", "partition_0", "partition_0"]
-        assert chunks[0].global_indexes == [0, 4, 8]
+        assert chunks[0].global_indexes == [10, 14, 18]
         assert len(chunks[1]) == 3
         assert chunks[1].partition_ids == ["partition_1", "partition_1", "partition_1"]
-        assert chunks[1].global_indexes == [1, 5, 9]
+        assert chunks[1].global_indexes == [11, 15, 19]
         assert len(chunks[2]) == 2
         assert chunks[2].partition_ids == ["partition_2", "partition_2"]
-        assert chunks[2].global_indexes == [2, 6]
+        assert chunks[2].global_indexes == [12, 16]
         assert len(chunks[3]) == 2
         assert chunks[3].partition_ids == ["partition_3", "partition_3"]
-        assert chunks[3].global_indexes == [3, 7]
+        assert chunks[3].global_indexes == [13, 17]
 
         # validate _custom_backend_meta is chunked
-        assert 0 in chunks[0].custom_meta
-        assert 4 in chunks[0].custom_meta
-        assert 8 in chunks[0].custom_meta
-        assert 1 not in chunks[0].custom_meta
-        assert 1 in chunks[1].custom_meta
+        assert 10 in chunks[0].custom_meta
+        assert 14 in chunks[0].custom_meta
+        assert 18 in chunks[0].custom_meta
+        assert 11 not in chunks[0].custom_meta
+        assert 11 in chunks[1].custom_meta
 
         # validate _custom_backend_meta is chunked
-        assert 0 in chunks[0]._custom_backend_meta
-        assert 4 in chunks[0]._custom_backend_meta
-        assert 8 in chunks[0]._custom_backend_meta
-        assert 1 not in chunks[0]._custom_backend_meta
-        assert 1 in chunks[1]._custom_backend_meta
+        assert 10 in chunks[0]._custom_backend_meta
+        assert 14 in chunks[0]._custom_backend_meta
+        assert 18 in chunks[0]._custom_backend_meta
+        assert 11 not in chunks[0]._custom_backend_meta
+        assert 11 in chunks[1]._custom_backend_meta
 
     def test_batch_meta_init_validation_error_different_field_names(self):
         """Example: Init validation catches samples with different field names."""

transfer_queue/metadata.py

@@ -409,13 +409,10 @@ def select_samples(self, indexes: list[int]) -> "BatchMeta":
 
         selected_samples = [self.samples[i] for i in indexes]
 
-        selected_custom_meta = {
-            i: self.custom_meta[self.global_indexes[i]] for i in indexes if self.global_indexes[i] in self.custom_meta
-        }
+        global_indexes = [self.global_indexes[i] for i in indexes]
+        selected_custom_meta = {i: self.custom_meta[i] for i in global_indexes if i in self.custom_meta}
         selected_custom_backend_meta = {
-            i: self._custom_backend_meta[self.global_indexes[i]]
-            for i in indexes
-            if self.global_indexes[i] in self._custom_backend_meta
+            i: self._custom_backend_meta[i] for i in global_indexes if i in self._custom_backend_meta
         }
 
         # construct new BatchMeta instance
@@ -470,11 +467,22 @@ def __getitem__(self, item):
         if isinstance(item, int | np.integer):
             sample_meta = self.samples[item] if self.samples else []
             global_idx = self.global_indexes[item]
+
+            if global_idx in self.custom_meta:
+                custom_meta = {global_idx: self.custom_meta[global_idx]}
+            else:
+                custom_meta = {}
+
+            if global_idx in self._custom_backend_meta:
+                custom_backend_meta = {global_idx: self._custom_backend_meta[global_idx]}
+            else:
+                custom_backend_meta = {}
+
             return BatchMeta(
                 samples=[sample_meta],
                 extra_info=self.extra_info,
-                custom_meta={global_idx: self.custom_meta[global_idx]},
-                _custom_backend_meta={global_idx: self._custom_backend_meta[global_idx]},
+                custom_meta=custom_meta,
+                _custom_backend_meta=custom_backend_meta,
             )
         else:
             raise TypeError(f"Indexing with {type(item)} is not supported now!")
@@ -533,11 +541,11 @@ def chunk_by_partition(
             List of smaller BatchMeta chunks, each chunk has samples with identical partition_id
         """
 
-        grouped_global_indexes = defaultdict(list)
-        for partition_id, global_index in zip(self.partition_ids, self.global_indexes, strict=False):
-            grouped_global_indexes[partition_id].append(global_index)
+        grouped_indexes = defaultdict(list)
+        for partition_id, indexes in zip(self.partition_ids, range(self.size), strict=False):
+            grouped_indexes[partition_id].append(indexes)
 
-        chunk_list = [self.select_samples(global_indices) for global_indices in grouped_global_indexes.values()]
+        chunk_list = [self.select_samples(idx) for idx in grouped_indexes.values()]
 
         return chunk_list