diff --git a/docs/api/data/pyhealth.data.Patient.rst b/docs/api/data/pyhealth.data.Patient.rst index 9965078db..24b38adb7 100644 --- a/docs/api/data/pyhealth.data.Patient.rst +++ b/docs/api/data/pyhealth.data.Patient.rst @@ -5,6 +5,17 @@ Another basic data structure in the package. A Patient is a collection of Visit the current patients. It contains all necessary attributes of a patient, such as ethnicity, mortality status, gender, etc. It can support various healthcare tasks. +Event order +----------- + +A patient's events are ordered by timestamp. Events that share a timestamp keep +the order of their rows in the source file: all diagnoses of a MIMIC admission +carry its discharge time, for example, so they come out in ``seq_num`` +(priority) order. Dataset loading records each row's position in a +``_source_row`` column of the cached event table, so the order is the same on +every cache build. Caches written before this column existed are rebuilt +automatically. + .. autoclass:: pyhealth.data.Patient :members: :undoc-members: diff --git a/examples/event_tie_order.py b/examples/event_tie_order.py new file mode 100644 index 000000000..b7ff8cbee --- /dev/null +++ b/examples/event_tie_order.py @@ -0,0 +1,52 @@ +"""Events that share a timestamp keep their source-file order. + +Builds a tiny dataset whose diagnoses all share one timestamp (as MIMIC +diagnoses share their admission's discharge time), then shows that a patient's +diagnoses come out in the order of the source file, which is their priority +order, on every cache build. + +Run: + python examples/event_tie_order.py +""" + +import tempfile +from pathlib import Path + +from pyhealth.datasets import BaseDataset + +CONFIG = """version: "1.0" +tables: + diagnoses: + file_path: "diagnoses.csv" + patient_id: "patient_id" + timestamp: "dischtime" + attributes: ["seq_num", "code"] +""" + +CODES = ["428.0", "250.00", "401.9", "038.9", "486"] # priority order + + +def main() -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) / "data" + root.mkdir() + (root / "config.yaml").write_text(CONFIG) + rows = ["patient_id,dischtime,seq_num,code"] + [ + f"p1,2020-01-05 12:00:00,{i + 1},{code}" for i, code in enumerate(CODES) + ] + (root / "diagnoses.csv").write_text("\n".join(rows) + "\n") + + for build in range(1, 4): + dataset = BaseDataset( + root=str(root), + tables=["diagnoses"], + dataset_name="tie_order_demo", + config_path=str(root / "config.yaml"), + cache_dir=str(Path(tmp) / f"cache{build}"), # a fresh build each time + ) + events = dataset.get_patient("p1").get_events(event_type="diagnoses") + print(f"build {build}:", [(e.seq_num, e.code) for e in events]) + + +if __name__ == "__main__": + main() diff --git a/pyhealth/data/data.py b/pyhealth/data/data.py index abc1b6e6e..8b9e2ddc2 100644 --- a/pyhealth/data/data.py +++ b/pyhealth/data/data.py @@ -130,10 +130,32 @@ def __getattr__(self, key: str) -> any: class Patient: """Patient class representing a sequence of events. + Events are ordered by timestamp. Events that share a timestamp (for + example all diagnoses of a MIMIC admission, which carry its discharge time) + keep the order of their rows in the source file, using the + ``_source_row`` column that dataset loading adds, so the order is the same + on every cache build. + Attributes: patient_id (str): Unique patient identifier. data_source (pl.DataFrame): DataFrame containing all events, sorted by timestamp. event_type_partitions (Dict[str, pl.DataFrame]): Dictionary mapping event types to their respective DataFrame partitions. + + Examples: + >>> from datetime import datetime + >>> import polars as pl + >>> from pyhealth.data import Patient + >>> t = datetime(2020, 1, 1) + >>> events = pl.DataFrame({ + ... "patient_id": ["p1"] * 3, + ... "event_type": ["diagnoses"] * 3, + ... "timestamp": [t, t, t], + ... "_source_row": [2, 0, 1], + ... "diagnoses/code": ["c", "a", "b"], + ... }) + >>> patient = Patient(patient_id="p1", data_source=events) + >>> [e.code for e in patient.get_events(event_type="diagnoses")] + ['a', 'b', 'c'] """ def __init__(self, patient_id: str, data_source: pl.DataFrame) -> None: @@ -145,7 +167,16 @@ def __init__(self, patient_id: str, data_source: pl.DataFrame) -> None: data_source (pl.DataFrame): DataFrame containing all events. """ self.patient_id = patient_id - self.data_source = data_source.sort("timestamp") + # Events that share a timestamp are ordered by table and then by their + # row in the source file, so the order does not depend on how the cache + # was built. Older caches lack the source-row column; there the + # existing row order is kept (a stable sort). + if "_source_row" in data_source.columns: + self.data_source = data_source.sort( + ["timestamp", "event_type", "_source_row"], maintain_order=True + ) + else: + self.data_source = data_source.sort("timestamp", maintain_order=True) self.event_type_partitions = self.data_source.partition_by("event_type", maintain_order=True, as_dict=True) def _filter_by_time_range_regular(self, df: pl.DataFrame, start: Optional[datetime], end: Optional[datetime]) -> pl.DataFrame: diff --git a/pyhealth/datasets/base_dataset.py b/pyhealth/datasets/base_dataset.py index fa71fabce..ff8f92d33 100644 --- a/pyhealth/datasets/base_dataset.py +++ b/pyhealth/datasets/base_dataset.py @@ -131,6 +131,23 @@ def path_exists(path: str) -> bool: return Path(path).exists() +# Column holding each event's row number in its source table (see load_table). +SOURCE_ROW = "_source_row" +# Layout of the cached event table; part of the dataset cache key, so caches +# written with an older layout are rebuilt. 2: adds SOURCE_ROW. +EVENT_FORMAT = 2 + + +def _with_source_row(df: dd.DataFrame) -> dd.DataFrame: + """Adds ``SOURCE_ROW``: 0, 1, 2, ... in the frame's current row order. + + A cumulative sum runs across partitions in order, unlike a per-partition + ``reset_index``. + """ + df = df.assign(**{SOURCE_ROW: 1}) + return df.assign(**{SOURCE_ROW: (df[SOURCE_ROW].cumsum() - 1).astype("int64")}) + + def _source_fingerprint(path: str) -> list: """Identifies the current version of a table source for the cache key. @@ -504,6 +521,7 @@ def _init_cache_dir(self, cache_dir: str | Path | None) -> Path: "tables": sorted(self.tables), "dataset_name": self.dataset_name, "dev": self.dev, + "event_format": EVENT_FORMAT, } if self.config is not None: # Editing the YAML or rewriting a source file must not reuse a stale @@ -801,6 +819,10 @@ def load_table(self, table_name: str) -> dd.DataFrame: # Convert column names to lowercase before calling preprocess_func df = df.rename(columns=str.lower) + # Number rows in source-file order. Patient uses this to order events + # that share a timestamp, so the order is the same on every cache build + # (the parallel build below does not preserve row order). + df = _with_source_row(df) # Check if there is a preprocessing function for this table preprocess_func: Optional[Callable[[nw.LazyFrame], nw.LazyFrame]] @@ -810,6 +832,8 @@ def load_table(self, table_name: str) -> dd.DataFrame: f"Preprocessing table: {table_name} with {preprocess_func.__name__}" ) df = preprocess_func(nw.from_native(df)).to_native() # type: ignore + if SOURCE_ROW not in df.columns: + df = _with_source_row(df) # Handle joins for join_cfg in table_cfg.join: @@ -880,7 +904,7 @@ def load_table(self, table_name: str) -> dd.DataFrame: df: dd.DataFrame = df.rename(columns=rename_attr) attr_cols = [rename_attr[attr.lower()] for attr in attribute_cols] - final_cols = ["patient_id", "event_type", "timestamp"] + attr_cols + final_cols = ["patient_id", "event_type", "timestamp", SOURCE_ROW] + attr_cols event_frame = df[final_cols] return event_frame diff --git a/tests/core/test_event_tie_order.py b/tests/core/test_event_tie_order.py new file mode 100644 index 000000000..155a763d2 --- /dev/null +++ b/tests/core/test_event_tie_order.py @@ -0,0 +1,140 @@ +"""Events that share a timestamp must come out in source-file order. + +The cached event table is built by a parallel sort on ``patient_id`` only, so +the order of a patient's rows differed between cache builds, and ``Patient`` +sorted by timestamp alone with an unstable sort. Events with equal timestamps +(all diagnoses of a MIMIC admission, same-day prescriptions, admissions that +start at the same time) therefore came out in a different order on different +builds, which changed task outputs (for example which admission counts as the +last one) and lost the ``seq_num`` priority order of diagnoses. +""" + +import random +import tempfile +import unittest +from datetime import datetime +from pathlib import Path +from unittest import mock + +import polars as pl + +from pyhealth.data import Patient +from pyhealth.datasets import BaseDataset +from pyhealth.datasets.base_dataset import EVENT_FORMAT, SOURCE_ROW + +CONFIG = """version: "1.0" +tables: + admissions: + file_path: "admissions.csv" + patient_id: "patient_id" + timestamp: "admittime" + attributes: + - "hadm_id" + diagnoses: + file_path: "diagnoses.csv" + patient_id: "patient_id" + timestamp: "time" + attributes: + - "hadm_id" + - "code" +""" + +# Codes in priority order, deliberately not sorted, and all at one timestamp. +CODES = ["428.0", "250.00", "401.9", "038.9", "486", "584.9", "276.1", "285.9"] + + +def _write_data(root: Path) -> None: + (root / "config.yaml").write_text(CONFIG) + (root / "admissions.csv").write_text( + "patient_id,admittime,hadm_id\n" + # Two admissions of p1 start at the same time. + "p1,2020-01-01 12:00:00,h2\n" + "p1,2020-01-01 12:00:00,h1\n" + "p1,2020-03-01 12:00:00,h3\n" + ) + rows = ["patient_id,time,hadm_id,code"] + rows += [f"p1,2020-01-05 00:00:00,h1,{code}" for code in CODES] + (root / "diagnoses.csv").write_text("\n".join(rows) + "\n") + + +def _events(n: int) -> pl.DataFrame: + t = datetime(2020, 1, 1) # noqa: DTZ001 - event timestamps are naive + return pl.DataFrame( + { + "patient_id": ["p1"] * n, + "event_type": ["diagnoses"] * n, + "timestamp": [t] * n, + SOURCE_ROW: list(range(n)), + "diagnoses/code": [f"c{i:03d}" for i in range(n)], + } + ) + + +class TestPatientTieOrder(unittest.TestCase): + def test_tied_events_follow_source_row_whatever_the_input_order(self): + events = _events(200) + expected = events["diagnoses/code"].to_list() + for seed in range(5): + rows = list(range(events.height)) + random.Random(seed).shuffle(rows) + patient = Patient("p1", events[rows]) + got = [e.code for e in patient.get_events(event_type="diagnoses")] + self.assertEqual(got, expected) + + def test_caches_without_source_row_keep_their_row_order(self): + events = _events(50).drop(SOURCE_ROW).reverse() + patient = Patient("p1", events) + got = [e.code for e in patient.get_events(event_type="diagnoses")] + self.assertEqual(got, events["diagnoses/code"].to_list()) + + def test_source_row_is_not_an_event_attribute(self): + patient = Patient("p1", _events(3)) + event = patient.get_events(event_type="diagnoses")[0] + self.assertNotIn(SOURCE_ROW, event.attr_dict) + + +class TestDatasetTieOrder(unittest.TestCase): + def setUp(self): + tmp = tempfile.TemporaryDirectory() + self.addCleanup(tmp.cleanup) + self.root = Path(tmp.name) / "data" + self.root.mkdir() + _write_data(self.root) + self.cache = Path(tmp.name) / "cache" + + def _dataset(self) -> BaseDataset: + return BaseDataset( + root=str(self.root), + tables=["admissions", "diagnoses"], + dataset_name="TieOrderDataset", + config_path=str(self.config_path), + cache_dir=str(self.cache), + ) + + @property + def config_path(self) -> Path: + return self.root / "config.yaml" + + def test_source_order_survives_the_cache_build(self): + patient = self._dataset().get_patient("p1") + admissions = [e.hadm_id for e in patient.get_events(event_type="admissions")] + self.assertEqual(admissions, ["h2", "h1", "h3"]) + codes = [e.code for e in patient.get_events(event_type="diagnoses")] + self.assertEqual(codes, CODES) + + def test_event_table_has_source_row(self): + columns = self._dataset().global_event_df.collect_schema().names() + self.assertIn(SOURCE_ROW, columns) + + def test_cache_key_includes_event_format(self): + # Caches written before SOURCE_ROW existed must not be reused. + current = self._dataset().cache_dir + self.assertEqual(self._dataset().cache_dir, current) + with mock.patch( + "pyhealth.datasets.base_dataset.EVENT_FORMAT", EVENT_FORMAT - 1 + ): + self.assertNotEqual(self._dataset().cache_dir, current) + + +if __name__ == "__main__": + unittest.main()