Key the dataset cache on the config and source files

docs/api/datasets.rst

@@ -31,6 +31,13 @@ compact ``global_event_df.parquet`` cache to disk. On subsequent runs with
 the same configuration it reads from cache rather than re-parsing the source
 files, so startup is fast.
 
+"The same configuration" means the same root, tables, dataset name and dev flag,
+the same YAML config, and unchanged source files: the cache key includes a hash of
+the config and the size and modification time of every source file the requested
+tables read. Editing the YAML or rewriting a source file (even at the same path)
+therefore builds a fresh cache instead of silently reusing a stale one. Old cache
+folders are left in place; delete them to reclaim disk space.
+
 The result is a :class:`~pyhealth.datasets.BaseDataset` — a structured
 patient→event tree. It is different from a PyTorch Dataset: it has no integer
 length and you cannot index into it with ``dataset[i]``. Think of it as a

examples/dataset_cache_refresh.py

@@ -0,0 +1,70 @@
+"""The dataset cache is rebuilt when a source file or the YAML config changes.
+
+PyHealth caches each dataset's event table. The cache key includes a hash of
+the YAML config and the size and modification time of every source file, so
+regenerating a source file at the same path (a common step in an analysis
+pipeline) gives a fresh cache instead of silently reusing stale data.
+
+Runs in a few seconds on CPU with a tiny synthetic table; no download.
+
+Usage:
+    python examples/dataset_cache_refresh.py
+"""
+
+import tempfile
+from pathlib import Path
+
+from pyhealth.datasets import BaseDataset
+
+CONFIG = """version: "1.0"
+tables:
+  visits:
+    file_path: "visits.csv"
+    patient_id: "patient_id"
+    timestamp: "visit_time"
+    attributes:
+      - "diagnosis"
+"""
+
+
+def write_visits(path: Path, patients: list[str]) -> None:
+    rows = ["patient_id,visit_time,diagnosis"]
+    rows += [f"{p},2024-01-15 09:00:00,J45" for p in patients]
+    path.write_text("\n".join(rows) + "\n")
+
+
+def load(root: Path, cache: Path) -> BaseDataset:
+    return BaseDataset(
+        root=str(root),
+        tables=["visits"],
+        dataset_name="cache_refresh_demo",
+        config_path=str(root / "config.yaml"),
+        cache_dir=str(cache),
+    )
+
+
+def n_patients(dataset: BaseDataset) -> int:
+    return dataset.global_event_df.select("patient_id").unique().collect().height
+
+
+def main():
+    with tempfile.TemporaryDirectory() as tmp:
+        root, cache = Path(tmp) / "data", Path(tmp) / "cache"
+        root.mkdir()
+        (root / "config.yaml").write_text(CONFIG)
+
+        write_visits(root / "visits.csv", ["p1", "p2"])
+        first = load(root, cache)
+        print("v1:", n_patients(first), "patients, cache", first.cache_dir.name)
+
+        # Regenerate the source file at the same path, e.g. after a cohort update.
+        write_visits(root / "visits.csv", ["p1", "p2", "p3", "p4"])
+        second = load(root, cache)
+        print("v2:", n_patients(second), "patients, cache", second.cache_dir.name)
+
+        again = load(root, cache)
+        print("unchanged rerun reuses cache:", again.cache_dir == second.cache_dir)
+
+
+if __name__ == "__main__":
+    main()

pyhealth/datasets/base_dataset.py

@@ -9,6 +9,8 @@
 from urllib.parse import urlparse, urlunparse
 from urllib.request import urlretrieve
 import json
+import glob
+import hashlib
 import uuid
 import platformdirs
 import multiprocessing
@@ -87,6 +89,41 @@ def path_exists(path: str) -> bool:
         return Path(path).exists()
 
 
+def _source_fingerprint(path: str) -> list:
+    """Identifies the current version of a table source for the cache key.
+
+    Returns ``[file, size, mtime_ns]`` for every local file behind ``path``:
+    the file itself and its ``.csv``/``.csv.gz`` (or ``.tsv``) twin, the files
+    matching a glob, or every file under a directory. A rewrite of any of them
+    therefore changes the dataset's cache key, like ``make``. Contents are not
+    hashed, so large sources stay cheap to check. URLs are identified by the
+    URL alone.
+
+    Args:
+        path (str): The table source path, as resolved from the config.
+
+    Returns:
+        list: ``[[file, size, mtime_ns], ...]``, or ``[path]`` for a URL.
+    """
+    if is_url(path):
+        return [path]
+    candidates = [path, path[:-3] if path.endswith(".gz") else f"{path}.gz"]
+    files: list[str] = []
+    for candidate in candidates:
+        source = Path(candidate)
+        if glob.has_magic(candidate):
+            files += sorted(glob.glob(candidate, recursive=True))
+        elif source.is_dir():
+            files += sorted(str(f) for f in source.rglob("*") if f.is_file())
+        elif source.is_file():
+            files.append(candidate)
+    fingerprint = []
+    for file in files:
+        stat = os.stat(file)
+        fingerprint.append([file, stat.st_size, stat.st_mtime_ns])
+    return fingerprint
+
+
 def _csv_tsv_gz_path(path: str) -> str:
     """
     Get the path to the file, trying the original path first, then the alternative path
@@ -375,6 +412,11 @@ def __init__(
     def _init_cache_dir(self, cache_dir: str | Path | None) -> Path:
         """Returns the cache directory path.
 
+        The cache key covers the dataset's root, tables, name and dev flag and,
+        for config-based datasets, a hash of the YAML config plus the size and
+        modification time of every source file the requested tables read. Editing
+        the config or rewriting a source file therefore builds a fresh cache.
+
         The cache directory is determined by the type of ``cache_dir`` passed
         to ``__init__``:
 
@@ -395,15 +437,26 @@ def _init_cache_dir(self, cache_dir: str | Path | None) -> Path:
         Returns:
             Path: The resolved cache directory path.
         """
-        id_str = json.dumps(
-            {
-                "root": str(self.root),
-                "tables": sorted(self.tables),
-                "dataset_name": self.dataset_name,
-                "dev": self.dev,
-            },
-            sort_keys=True,
-        )
+        key = {
+            "root": str(self.root),
+            "tables": sorted(self.tables),
+            "dataset_name": self.dataset_name,
+            "dev": self.dev,
+        }
+        if self.config is not None:
+            # Editing the YAML or rewriting a source file must not reuse a stale
+            # cache, so the key also covers the config and the source files.
+            config_json = json.dumps(self.config.model_dump(mode="json"), sort_keys=True)
+            key["config"] = hashlib.sha256(config_json.encode()).hexdigest()
+            key["sources"] = {
+                table: [
+                    _source_fingerprint(clean_path(f"{self.root}/{file_path}"))
+                    for file_path in [cfg.file_path, *(j.file_path for j in cfg.join)]
+                ]
+                for table in sorted(self.tables)
+                if (cfg := self.config.tables.get(table)) is not None
+            }
+        id_str = json.dumps(key, sort_keys=True)
 
         id = str(uuid.uuid5(uuid.NAMESPACE_DNS, id_str))
 

tests/core/test_dataset_cache_key.py

@@ -0,0 +1,91 @@
+"""The dataset cache key must change when the source data or config changes.
+
+Before, the key covered only {root, tables, dataset_name, dev}: rewriting a
+source file at the same path, or editing the YAML config, silently reused the
+stale cached event table.
+"""
+
+import os
+import tempfile
+import time
+import unittest
+from pathlib import Path
+
+from pyhealth.datasets import BaseDataset
+
+CONFIG = """version: "1.0"
+tables:
+  events:
+    file_path: "events.csv"
+    patient_id: "patient_id"
+    timestamp: "time"
+    attributes:
+      - "code"
+"""
+
+
+def _write_csv(path: Path, n_patients: int) -> None:
+    rows = ["patient_id,time,code"] + [
+        f"p{i},2020-01-0{1 + i % 9} 00:00:00,c{i}" for i in range(n_patients)
+    ]
+    path.write_text("\n".join(rows) + "\n")
+
+
+class TestDatasetCacheKey(unittest.TestCase):
+    def setUp(self):
+        tmp = tempfile.TemporaryDirectory()
+        self.addCleanup(tmp.cleanup)
+        self.root = Path(tmp.name) / "data"
+        self.cache = Path(tmp.name) / "cache"
+        self.root.mkdir()
+        self.config = self.root / "config.yaml"
+        self.config.write_text(CONFIG)
+        self.csv = self.root / "events.csv"
+        _write_csv(self.csv, n_patients=3)
+
+    def _dataset(self) -> BaseDataset:
+        return BaseDataset(
+            root=str(self.root),
+            tables=["events"],
+            dataset_name="CacheKeyDataset",
+            config_path=str(self.config),
+            cache_dir=str(self.cache),
+        )
+
+    def _n_patients(self, dataset: BaseDataset) -> int:
+        return dataset.global_event_df.select("patient_id").unique().collect().height
+
+    def test_rewritten_source_file_is_reloaded(self):
+        first = self._dataset()
+        self.assertEqual(self._n_patients(first), 3)
+
+        time.sleep(0.01)  # make sure the modification time moves
+        _write_csv(self.csv, n_patients=5)  # same path, new content
+
+        second = self._dataset()
+        self.assertNotEqual(first.cache_dir, second.cache_dir)
+        self.assertEqual(self._n_patients(second), 5)
+
+    def test_unchanged_source_reuses_the_cache(self):
+        first = self._dataset()
+        _ = first.global_event_df
+        second = self._dataset()
+        self.assertEqual(first.cache_dir, second.cache_dir)
+
+    def test_edited_config_changes_the_cache(self):
+        first = self._dataset()
+        self.config.write_text(CONFIG.replace('      - "code"\n', '      - "code"\n      - "time"\n'))
+        second = self._dataset()
+        self.assertNotEqual(first.cache_dir, second.cache_dir)
+
+    def test_touching_mtime_only_changes_the_cache(self):
+        # Size + modification time is the signal (like make); content is not hashed.
+        first = self._dataset()
+        stat = self.csv.stat()
+        os.utime(self.csv, ns=(stat.st_atime_ns, stat.st_mtime_ns + 10**9))
+        second = self._dataset()
+        self.assertNotEqual(first.cache_dir, second.cache_dir)
+
+
+if __name__ == "__main__":
+    unittest.main()