Fix SampleDataset cache cleanup

docs/api/datasets.rst

@@ -82,6 +82,9 @@ future runs without re-fitting.
 Calling ``samples.subset(...)`` rebuilds both lookups with indices local to the
 new dataset, so they remain valid after repeated splitting.
 
+Calling ``close()`` preserves reusable task caches. Disk-backed datasets created
+by ``create_sample_dataset()`` own a temporary directory and remove it on close.
+
 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

@@ -14,13 +14,15 @@ def main():
         samples=samples,
         input_schema={"feature": "raw"},
         output_schema={"label": "raw"},
+        in_memory=False,
     )
 
     subset = dataset.subset([3, 0, 2])
     nested_subset = subset.subset([2, 0])
 
     print(subset.patient_to_index)
     print(nested_subset.patient_to_index)
+    dataset.close()
 
 
 if __name__ == "__main__":

pyhealth/datasets/sample_dataset.py

@@ -313,6 +313,7 @@ def __init__(
         path: str,
         dataset_name: Optional[str] = None,
         task_name: Optional[str] = None,
+        delete_on_close: bool = False,
         **kwargs,
     ) -> None:
         """Initialize a SampleDataset pointing at a directory created by SampleBuilder.
@@ -322,6 +323,8 @@ def __init__(
                 `SampleBuilder.save` and associated pickled sample files.
             dataset_name: Optional human-friendly dataset name.
             task_name: Optional human-friendly task name.
+            delete_on_close: Whether ``close()`` should delete ``path``. This is
+                reserved for temporary datasets owned by PyHealth.
             **kwargs: Extra keyword arguments forwarded to
                 `litdata.StreamingDataset` (such as streaming options).
         """
@@ -330,6 +333,7 @@ def __init__(
         self.path = path
         self.dataset_name = "" if dataset_name is None else dataset_name
         self.task_name = "" if task_name is None else task_name
+        self._delete_on_close = delete_on_close
 
         with open(f"{path}/schema.pkl", "rb") as f:
             metadata = pickle.load(f)
@@ -442,13 +446,18 @@ def subset(self, indices: Union[Sequence[int], slice]) -> "SampleDataset":
         new_dataset.record_to_index = _remap_index_mapping(
             self.record_to_index, indices
         )
+        new_dataset._delete_on_close = False
         new_dataset.reset()
 
         return new_dataset
 
     def close(self) -> None:
         """Cleans up any temporary directories used by the dataset."""
-        if self.input_dir.path is not None and Path(self.input_dir.path).exists():
+        if (
+            self._delete_on_close
+            and self.input_dir.path is not None
+            and Path(self.input_dir.path).exists()
+        ):
             shutil.rmtree(self.input_dir.path)
 
     # --------------------------------------------------------------
@@ -632,6 +641,15 @@ def create_sample_dataset(
     Returns:
         An instance of `SampleDataset` loaded from the temporary directory
         containing the optimized, chunked samples and `schema.pkl` metadata.
+
+    Examples:
+        >>> dataset = create_sample_dataset(
+        ...     samples=[{"patient_id": "p1", "feature": 1, "label": 0}],
+        ...     input_schema={"feature": "raw"},
+        ...     output_schema={"label": "raw"},
+        ... )
+        >>> len(dataset)
+        1
     """
     if in_memory:
         return InMemorySampleDataset(
@@ -666,4 +684,5 @@ def create_sample_dataset(
             path=str(path),
             dataset_name=dataset_name,
             task_name=task_name,
+            delete_on_close=True,
         )

tests/core/test_sample_dataset.py

@@ -1,7 +1,8 @@
 import unittest
-import pickle
-import random
-from pyhealth.datasets.sample_dataset import create_sample_dataset
+from pathlib import Path
+
+from pyhealth.datasets.sample_dataset import SampleDataset, create_sample_dataset
+
 
 class TestSampleDatasetParity(unittest.TestCase):
     def setUp(self):
@@ -143,5 +144,37 @@ def test_set_shuffle(self):
             self.assertEqual(items_disk_ordered[i]["feature"], i)
             self.assertEqual(items_mem_ordered[i]["feature"], i)
 
+    def test_close_preserves_caller_owned_directory(self):
+        owner, _ = self._get_datasets()
+        path = Path(owner.path)
+        dataset = SampleDataset(path=str(path))
+
+        try:
+            dataset.close()
+            self.assertTrue(path.exists())
+            reopened = SampleDataset(path=str(path))
+            self.assertEqual(reopened[0]["feature"], 0)
+        finally:
+            owner.close()
+
+    def test_close_preserves_directory_shared_with_subset(self):
+        owner, _ = self._get_datasets()
+        path = Path(owner.path)
+        subset = owner.subset([0])
+
+        try:
+            subset.close()
+            self.assertTrue(path.exists())
+        finally:
+            owner.close()
+
+    def test_close_removes_owned_temporary_directory(self):
+        owner, _ = self._get_datasets()
+        path = Path(owner.path)
+
+        owner.close()
+
+        self.assertFalse(path.exists())
+
 if __name__ == "__main__":
     unittest.main()