diff --git a/docs/api/datasets.rst b/docs/api/datasets.rst index 007014a5b..0955b46d2 100644 --- a/docs/api/datasets.rst +++ b/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 diff --git a/examples/dataset_cache_refresh.py b/examples/dataset_cache_refresh.py new file mode 100644 index 000000000..7d1c65c40 --- /dev/null +++ b/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() diff --git a/pyhealth/datasets/base_dataset.py b/pyhealth/datasets/base_dataset.py index 8718eafa1..01314a9e4 100644 --- a/pyhealth/datasets/base_dataset.py +++ b/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)) diff --git a/tests/core/test_dataset_cache_key.py b/tests/core/test_dataset_cache_key.py new file mode 100644 index 000000000..17aff2f2e --- /dev/null +++ b/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()