diff --git a/pyhealth/datasets/base_dataset.py b/pyhealth/datasets/base_dataset.py index 3d449d579..88c930534 100644 --- a/pyhealth/datasets/base_dataset.py +++ b/pyhealth/datasets/base_dataset.py @@ -749,7 +749,13 @@ def load_table(self, table_name: str) -> dd.DataFrame: df: dd.DataFrame = df.assign(patient_id=df[patient_id_col].astype("string")) else: df: dd.DataFrame = df.reset_index(drop=True) - df: dd.DataFrame = df.assign(patient_id=df.index.astype("string")) + + # Dask applies reset_index independently across frames, + # meaning partitions share indexes, and therefore merge patients. + # The trick is to use a cumsum which acts globally across partitions. + df: dd.DataFrame = df.assign(_one=1) + df: dd.DataFrame = df.assign(patient_id=(df["_one"].cumsum() - 1).astype("string")) + df: dd.DataFrame = df.drop(columns="_one") df: dd.DataFrame = df.assign(event_type=table_name) diff --git a/pyhealth/datasets/fhir/base.py b/pyhealth/datasets/fhir/base.py index 517d062bb..67205c15c 100644 --- a/pyhealth/datasets/fhir/base.py +++ b/pyhealth/datasets/fhir/base.py @@ -383,7 +383,13 @@ def load_table(self, table_name: str) -> dd.DataFrame: df = df.assign(patient_id=df[table_cfg.patient_id].astype("string")) else: df = df.reset_index(drop=True) - df = df.assign(patient_id=df.index.astype("string")) + + # Dask applies reset_index independently across frames, + # meaning partitions share indexes, and therefore merge patients. + # The trick is to use a cumsum which acts globally across partitions. + df = df.assign(_one=1) + df = df.assign(patient_id=(df["_one"].cumsum() - 1).astype("string")) + df = df.drop(columns="_one") df = df.dropna(subset=["patient_id"]) df = df.assign(event_type=table_name) diff --git a/tests/core/test_base_dataset.py b/tests/core/test_base_dataset.py index 4f9bb1fda..11ee0843c 100644 --- a/tests/core/test_base_dataset.py +++ b/tests/core/test_base_dataset.py @@ -296,6 +296,37 @@ class ConcreteDataset(BaseDataset): self.assertTrue(pd.isna(pdf.iloc[2]["timestamp"])) self.assertEqual(pdf.iloc[2]["table1/val"], "v3") + def test_null_patient_id_unique_across_partitions(self): + """Regression test for the Dask per-partition reset_index bug. + + A table config with ``patient_id: null`` must give every row a + globally unique ``patient_id``, even when the frame spans multiple + Dask partitions. The old code used ``df.reset_index(drop=True)`` + + ``df.index``, which Dask resets per partition, so rows from different + partitions collided onto the same id and were silently merged into + one "patient". + """ + from pyhealth.datasets.configs.config import DatasetConfig + + class _NullPatientDataset(BaseDataset): + def __init__(self, config, root="/tmp/x"): + self.config = config + self.root = root + + ds = _NullPatientDataset( + DatasetConfig( + version="1.0", + tables={"t": {"file_path": "t.csv", "attributes": ["value"]}}, + ) + ) + # 6 rows across 3 partitions; the buggy code produced ids [0,1,0,1,0,1]. + frame = dd.from_pandas(pd.DataFrame({"value": list(range(6))}), npartitions=3) + with patch.object(ds, "_scan_table", return_value=frame): + out = ds.load_table("t").compute() + self.assertEqual(len(out), 6) + self.assertEqual(out["patient_id"].nunique(), 6) + self.assertEqual(sorted(out["patient_id"].tolist()), [str(i) for i in range(6)]) + if __name__ == "__main__": unittest.main() diff --git a/tests/core/test_fhir_dataset.py b/tests/core/test_fhir_dataset.py index 1f8557f9a..b46d3b440 100644 --- a/tests/core/test_fhir_dataset.py +++ b/tests/core/test_fhir_dataset.py @@ -813,5 +813,45 @@ def pre_filter(self, df: pl.LazyFrame) -> pl.LazyFrame: self.assertIn("http://loinc.org|789-0", vocab.token_to_id) + +class TestFHIRNullPatientIdUniqueAcrossPartitions(unittest.TestCase): + """Regression test: FHIRDataset.load_table must give every row a globally + unique patient_id for a ``patient_id: null`` table across multiple Dask + partitions (same per-partition reset_index bug as BaseDataset).""" + + def test_patient_ids_unique_across_partitions(self): + import tempfile + import pandas as pd + import dask.dataframe as dd + from unittest.mock import patch + from pyhealth.datasets.fhir.base import FHIRDataset + from pyhealth.datasets.configs.config import DatasetConfig + + class _NullPatientFHIR(FHIRDataset): + def __init__(self, config, prepared_dir): + self.config = config + self._prepared = Path(prepared_dir) + self.output_format = "csv" + + @property + def prepared_tables_dir(self): + return self._prepared + + with tempfile.TemporaryDirectory() as d: + (Path(d) / "t.csv").write_text("") # satisfy the path.exists() guard + config = DatasetConfig( + version="1.0", + tables={"t": {"file_path": "t.csv", "attributes": ["value"]}}, + ) + ds = _NullPatientFHIR(config, d) + frame = dd.from_pandas(pd.DataFrame({"value": list(range(6))}), npartitions=3) + with patch.object(ds, "_read_flat_table", return_value=frame): + out = ds.load_table("t").compute() + self.assertEqual(len(out), 6) + self.assertEqual(out["patient_id"].nunique(), 6) + self.assertEqual(sorted(out["patient_id"].tolist()), [str(i) for i in range(6)]) + + + if __name__ == "__main__": unittest.main()