Order same-timestamp events by source-file row, so set_task is the same on every cache build

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:

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()

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:

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

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()