Fix SampleDataset subset mappings

docs/api/datasets.rst

@@ -79,6 +79,9 @@ future runs without re-fitting.
 - ``samples.record_to_index`` — maps a visit/record ID to the sample indices
   for that visit.
 
+Calling ``samples.subset(...)`` rebuilds both lookups with indices local to the
+new dataset, so they remain valid after repeated splitting.
+
 For testing or small cohorts you can skip the disk step entirely using
 ``InMemorySampleDataset``, which holds all processed samples in RAM and is
 returned by default from ``create_sample_dataset()``.

examples/sample_dataset_subset.py

@@ -0,0 +1,27 @@
+"""Create reordered and nested SampleDataset subsets."""
+
+from pyhealth.datasets import create_sample_dataset
+
+
+def main():
+    samples = [
+        {"patient_id": "p1", "record_id": "r1", "feature": 1, "label": 0},
+        {"patient_id": "p2", "record_id": "r2", "feature": 2, "label": 1},
+        {"patient_id": "p1", "record_id": "r3", "feature": 3, "label": 0},
+        {"patient_id": "p3", "record_id": "r4", "feature": 4, "label": 1},
+    ]
+    dataset = create_sample_dataset(
+        samples=samples,
+        input_schema={"feature": "raw"},
+        output_schema={"label": "raw"},
+    )
+
+    subset = dataset.subset([3, 0, 2])
+    nested_subset = subset.subset([2, 0])
+
+    print(subset.patient_to_index)
+    print(nested_subset.patient_to_index)
+
+
+if __name__ == "__main__":
+    main()

pyhealth/datasets/sample_dataset.py

@@ -15,6 +15,22 @@
 from ..processors.base_processor import FeatureProcessor
 
 
+def _remap_index_mapping(
+    mapping: dict[str, list[int]], indices: Sequence[int]
+) -> dict[str, list[int]]:
+    key_by_index = {
+        index: key
+        for key, mapped_indices in mapping.items()
+        for index in mapped_indices
+    }
+    remapped: dict[str, list[int]] = {}
+    for new_index, old_index in enumerate(indices):
+        key = key_by_index.get(old_index)
+        if key is not None:
+            remapped.setdefault(key, []).append(new_index)
+    return remapped
+
+
 class SampleBuilder:
     """Fit feature processors and transform pickled samples without materializing a dataset.
 
@@ -275,6 +291,21 @@ class SampleDataset(litdata.StreamingDataset):
             sample indices associated with that record.
         dataset_name: Optional human friendly dataset name.
         task_name: Optional human friendly task name.
+
+    Examples:
+        >>> from pyhealth.datasets import create_sample_dataset
+        >>> dataset = create_sample_dataset(  # doctest: +SKIP
+        ...     samples=[
+        ...         {"patient_id": "p1", "feature": 1, "label": 0},
+        ...         {"patient_id": "p2", "feature": 2, "label": 1},
+        ...     ],
+        ...     input_schema={"feature": "raw"},
+        ...     output_schema={"label": "raw"},
+        ...     in_memory=False,
+        ... )
+        >>> dataset.subset([1]).patient_to_index  # doctest: +SKIP
+        {'p2': [0]}
+        >>> dataset.close()  # doctest: +SKIP
     """
 
     def __init__(
@@ -355,6 +386,8 @@ def subset(self, indices: Union[Sequence[int], slice]) -> "SampleDataset":
         if isinstance(indices, slice):
             indices = range(*indices.indices(dataset_length))
 
+        indices = list(indices)
+
         if any(idx < 0 or idx >= dataset_length for idx in indices):
             raise ValueError(
                 f"Subset indices must be in [0, {dataset_length - 1}] for the provided dataset."
@@ -403,6 +436,12 @@ def subset(self, indices: Union[Sequence[int], slice]) -> "SampleDataset":
 
         new_dataset.subsampled_files = new_subsampled_files
         new_dataset.region_of_interest = new_roi
+        new_dataset.patient_to_index = _remap_index_mapping(
+            self.patient_to_index, indices
+        )
+        new_dataset.record_to_index = _remap_index_mapping(
+            self.record_to_index, indices
+        )
         new_dataset.reset()
 
         return new_dataset
@@ -430,6 +469,16 @@ class InMemorySampleDataset(SampleDataset):
     for fast, repeated access to samples without disk I/O, at the cost of
     higher memory usage.
 
+    Examples:
+        >>> from pyhealth.datasets import create_sample_dataset
+        >>> dataset = create_sample_dataset(
+        ...     samples=[{"patient_id": "p1", "feature": 1, "label": 0}],
+        ...     input_schema={"feature": "raw"},
+        ...     output_schema={"label": "raw"},
+        ... )
+        >>> dataset[0]["patient_id"]
+        'p1'
+
     Note:
         This class is intended for testing and debugging purposes where
         dataset sizes are small enough to fit into memory.
@@ -522,12 +571,23 @@ def __iter__(self) -> Iterable[Dict[str, Any]]:  # type: ignore
 
     def subset(self, indices: Union[Sequence[int], slice]) -> SampleDataset:
         if isinstance(indices, slice):
+            subset_indices = list(range(*indices.indices(len(self))))
             samples = self._data[indices]
         else:
-            samples = [self._data[i] for i in indices]
+            raw_indices = list(indices)
+            samples = [self._data[i] for i in raw_indices]
+            subset_indices = [
+                index if index >= 0 else len(self) + index for index in raw_indices
+            ]
 
         new_dataset = copy.deepcopy(self)
         new_dataset._data = samples
+        new_dataset.patient_to_index = _remap_index_mapping(
+            self.patient_to_index, subset_indices
+        )
+        new_dataset.record_to_index = _remap_index_mapping(
+            self.record_to_index, subset_indices
+        )
         return new_dataset
 
     def close(self) -> None:

tests/core/test_sample_dataset.py

@@ -83,6 +83,30 @@ def test_subset_slice(self):
         for d, m in zip(list_disk, list_mem):
             self.assertEqual(d["feature"], m["feature"])
 
+    def test_subset_rebuilds_index_mappings(self):
+        ds_disk, ds_mem = self._get_datasets()
+
+        for dataset in (ds_disk, ds_mem):
+            subset = dataset.subset([5, 2, 15])
+            self.assertEqual(
+                subset.patient_to_index,
+                {"p5": [0], "p2": [1], "p15": [2]},
+            )
+            self.assertEqual(
+                subset.record_to_index,
+                {"r5": [0], "r2": [1], "r15": [2]},
+            )
+
+            nested_subset = subset.subset([2, 0])
+            self.assertEqual(
+                nested_subset.patient_to_index,
+                {"p15": [0], "p5": [1]},
+            )
+            self.assertEqual(
+                nested_subset.record_to_index,
+                {"r15": [0], "r5": [1]},
+            )
+
     def test_set_shuffle(self):
         ds_disk, ds_mem = self._get_datasets()