Add MEDS dataset support (MEDSDataset + typed Parquet scan path) by AxelNoun · Pull Request #1179 · sunlabuiuc/PyHealth

AxelNoun · 2026-07-15T19:42:50Z

Motivation

The Medical Event Data Standard (MEDS) is a minimal, event-based schema for machine learning over EHR data (MEDS Working Group, Medical Event Data Standard (MEDS): Facilitating Machine Learning for Health, ICLR 2024 Workshop on Learning from Time Series For Health, openreview:IsHy2ebjIG). MEDS datasets are distributed as sharded, natively typed Parquet files.

PyHealth currently cannot read Parquet sources: .parquet paths and directories fall through the CSV/TSV extension resolver and raise ValueError. Since the MEDS core table maps one-to-one onto PyHealth's canonical event frame (subject_id → patient_id, time → timestamp, code/numeric_value as attributes), MEDS support is almost purely declarative once Parquet can be scanned.

Related to #329 (Integration with MEDS).

Changes

1. BaseDataset: generic Parquet source support (~90 changed lines)

_scan_table: format router — .parquet/.pq files, glob patterns, and directories go to the new scanner; everything else falls back to the existing _scan_csv_tsv_gz, unchanged. Non-breaking by construction: all newly routed inputs previously raised.
_scan_parquet: typed scan (no all-string coercion — Parquet embeds its schema) with the same Dask kwargs as the cached-Parquet read of the CSV path. A directory is scanned recursively, which supports the sharded MEDS layout data/<split>/<shard>.parquet. Missing or Parquet-less sources raise FileNotFoundError.
load_table: a fast-path when the configured timestamp column is already datetime-typed — skips the string round-trip, preserves NaT (MEDS static events have null time), then applies the usual datetime64[ms] cast. The list-concatenation and string-parsing branches are untouched.
A >>> usage example was added to the BaseDataset docstring, as required by the PR rules once the class span is modified.

2. MEDSDataset + configs/meds.yaml

Declarative wrapper patterned on EHRShotDataset; an optional static subject_splits table mirrors the EHRShot splits precedent.
subset="train" | "tuning" | "held_out", resolved via metadata/subject_splits.parquet (with a directory-layout fallback); each subset gets its own processing cache. Linked in source docs and the API reference page to the MEDS schema documentation.
A construction-time schema guard reads Parquet footers only (no data, no Dask): a missing, non-timestamp, or timezone-aware time column raises TypeError. The MEDS reference schema is timezone-naive timestamp[us]; the guard also rejects date-like integer columns (e.g. 20240101) that the generic string fallback would otherwise parse silently.

3. Tests — tests/core/test_meds.py

Deterministic synthetic sharded fixtures: nested split directories, static null-time events, µs→ms precision, subset filtering, distinct caches per subset, schema-guard violations. CI needs no external data.
Demo-backed smoke tests are skip-gated behind MEDS_DEMO_ROOT (public MIMIC-IV demo data in MEDS, PhysioNet, ODbL v1.0).
The CSV path shared by existing datasets is regression-covered by the current suite.

4. Docs & example

docs/api/datasets/pyhealth.datasets.MEDSDataset.rst
examples/meds_demo.py: load/split peek, then end-to-end InHospitalMortalityMEDS → RNN → Trainer on the public demo (238 stays, 80/10/10 split). Full-stay code sequences are tail-truncated to 256 events in a demo-only wrapper so the script finishes on CPU (~2.5 min, mostly set_task I/O); not a benchmark.

5. InHospitalMortalityMEDS (MEDS-native task)

One sample per completed stay: group events on hadm_id from HOSPITAL_ADMISSION//* / HOSPITAL_DISCHARGE//*, label = HOSPITAL_DISCHARGE//DIED, features = ordered codes in half-open [admit, prediction_time) excluding discharge and MEDS_DEATH (anti-leakage). Default window full_stay; opt-in first_hours.
hadm_id is not core MEDS, so it stays out of configs/meds.yaml and is exposed via bundled configs/meds_with_hadm.yaml (absent attributes raise in load_table).
Tests: synthetic leakage/window/label coverage + optional demo smoke; examples/verify_meds_mortality.py prints cohort stats through the real pipeline.
API docs: docs/api/tasks/pyhealth.tasks.InHospitalMortalityMEDS.rst + toctree entry in docs/api/tasks.rst.

Testing

python -m pytest tests/core/test_meds.py tests/core/test_in_hospital_mortality_meds.py -v: 24 passed (incl. schema-guard subtests + mortality synthetic/demo smoke) — Windows, Python 3.12.
Local smoke on the public demo: 100 subjects / 803,992 events; canonical dtypes verified (patient_id string, timestamp datetime64[ms]).
Cohort via real task pipeline (examples/verify_meds_mortality.py --root <demo>): 12 positive / 238 stays (rate 0.0504); set_task sample count 238.
examples/meds_demo.py (optional, requires demo download): completes on CPU in ~2.5 min; training epoch ~12 s on truncated sequences.

How to verify

pip install -e ".[dev]"
python -m pytest tests/core/test_meds.py tests/core/test_in_hospital_mortality_meds.py -v

# Optional: demo-backed tests + examples on the public PhysioNet demo (ODbL)
wget -r -N -c -np https://physionet.org/files/mimic-iv-demo-meds/0.0.1/
MEDS_DEMO_ROOT=physionet.org/files/mimic-iv-demo-meds/0.0.1 \
    python -m pytest tests/core/test_meds.py tests/core/test_in_hospital_mortality_meds.py -v
python examples/meds_demo.py --root physionet.org/files/mimic-iv-demo-meds/0.0.1
python examples/verify_meds_mortality.py --root physionet.org/files/mimic-iv-demo-meds/0.0.1

Notes for review

Scope: the BaseDataset Parquet path is deliberately generic (reusable by future Parquet-distributed datasets). Happy to split it into its own PR if you prefer smaller diffs.
subset API: small and tested, but easy to strip if you want the first cut minimal.
Demo-backed tests: currently skip-gated so the repo carries no new data. Alternatively I can commit a reduced ODbL slice (+ LICENSE/NOTICE) under test-resources/, following the eICU/MIMIC demo precedent — your call.
hadm_id config: stay-aware tasks opt into meds_with_hadm.yaml; the default MEDS config stays core-schema-only so generic MEDS exports without hadm_id keep working.

Route .parquet/.pq files, globs, and directories through a typed _scan_parquet scanner; keep CSV/TSV(.gz) on the existing path. Add a datetime fast-path in load_table that skips the string round-trip and casts to datetime64[ms], preserving NaT for static events.

Declarative YAML wrapper over the shared Parquet scan path, with split_source subset selection (metadata or directory layout), distinct processing caches per subset, and a construction-time Parquet footer schema guard that rejects missing, non-timestamp, or timezone-aware time columns.

Deterministic sharded Parquet fixtures cover nested splits, subset filtering, cache isolation, set_task smoke, and construction-time schema-guard TypeErrors. Demo smoke stays skip-gated behind MEDS_DEMO_ROOT / test-resources/meds_demo (gitignored).

Document MEDSDataset in the API reference and add an end-to-end examples/meds_demo.py against the public PhysioNet MIMIC-IV MEDS demo.

MEDS-native in-hospital mortality task: one sample per completed stay, reconstructed by joining HOSPITAL_ADMISSION/HOSPITAL_DISCHARGE events on hadm_id. Half-open [admit, prediction_time) observation window (full_stay default; first_hours early-warning variant), label from the HOSPITAL_DISCHARGE//DIED discharge code. Discharge and MEDS_DEATH events are excluded from features to prevent label leakage. hadm_id is dataset-specific (not part of the core MEDS schema), so it is exposed via a bundled configs/meds_with_hadm.yaml rather than the default config. Verified on the public MIMIC-IV demo in MEDS: 12 positive / 238 stays (rate 0.0504); set_task sample count 238. - pyhealth/tasks/in_hospital_mortality_meds.py (+ __init__ export) - pyhealth/datasets/configs/meds_with_hadm.yaml - tests/core/test_in_hospital_mortality_meds.py - examples/verify_meds_mortality.py - docs/api/tasks/pyhealth.tasks.InHospitalMortalityMEDS.rst (+ tasks.rst toctree) Co-authored-by: Cursor <cursoragent@cursor.com>

Satisfies ruff F401 on the newly added __init__ line under tools/check_pr_rules scoped lint (same pattern as eegbci). Co-authored-by: Cursor <cursoragent@cursor.com>

jhnwu3

Sorry for the late reply! I've got a couple small nitpicks, but yeah this is basically how I would've done it!

I've got 4 small comments, thanks for the hard work!

jhnwu3 · 2026-07-26T21:29:55Z

+MEDS distributes event data as *typed*, sharded Parquet, already flattened to
+one row per measurement -- ``(subject_id, time, code, numeric_value, ...)`` --
+plus a canonical subject-to-split mapping at
+``metadata/subject_splits.parquet``. This maps almost one-to-one onto


Is it possible to reference this split mapping here to a MEDS doc link? Just asking to make sure users who read the docs know about this.

Added: linked the MEDS schema docs next to the subject_splits.parquet mention.

Follow-up: also added the link to the rendered API page intro (docs/api/datasets/pyhealth.datasets.MEDSDataset.rst), in addition to the class docstring in meds.py.

jhnwu3 · 2026-07-26T21:32:49Z

+            earliest-admission / latest-discharge aggregation.
+        """
+        code = pl.col("meds/code")
+        admissions = (


ah, I see what you're doing here. Was going to ask for more clarification, but maybe "aggregate_stays" might be a better term. Reconstruction implies something like actual reconstruction loss, etc. here.

Agreed, reconstruction was misleading. Renamed to _group_stays; stays are grouped on hadm_id, not reconstructed. Prose updated to match.

jhnwu3 · 2026-07-26T21:34:01Z

+    @staticmethod
+    def summarize(samples: List[Dict[str, Any]]) -> Dict[str, Any]:
+        """Summary statistics of a generated sample set, for cohort reporting.
+


I'm curious what this is used for? For reference, the BaseDataset.set_task() -> SampleDataset which is a PyTorch data class, meaning it's not exactly a List[Dict] here.

It takes the raw List[Dict] from call (per-patient application), not the tokenized SampleDataset from set_task(). Callers (examples/verify_meds_mortality.py, tests) use it that way deliberately, so cohort stats stay inspectable on the string codes before processors are fitted. Clarified the docstring. Happy to move it out of the task class if you'd rather it live elsewhere.

Oh so literally .summarize() before set_task()?

Hmm, I think you might need to change how it works, because it'll be really slow for much larger datasets as it doesn't reuse the set_task() operations and it won't be memory efficient since it's not using our lazyloading?. I think it's easier to do it post-tokenization here, because then you can leverage the parallel threading to do the extraction first and then everything is already a tensor.

Fair point; the per-patient loop bypasses lazy loading and the parallel processing in set_task(), so it wouldn't scale past demo-sized data. Removed it, along with its call sites in the test and example. Happy to add a SampleDataset-based utility as a follow-up if cohort stats turn out to be worth having; it probably belongs outside the task class either way.

jhnwu3 · 2026-07-26T21:36:19Z

+    n_subset = len(subset.unique_patient_ids)
+    n_total = len(dataset.unique_patient_ids)
+    print(f"Subjects in subset '{args.subset}': {n_subset} / {n_total}")
+


Can we try training an RNN model at all on this? from pyhealth.models. That would help us showcase it works.

Added: set_task(InHospitalMortalityMEDS()) → pyhealth.models.RNN → Trainer, on the full demo cohort (238 stays, 80/10/10). Full-stay sequences run ~3k codes, so the demo truncates to the last 256 per stay via a local wrapper task API unchanged. ~2.5 min on CPU, mostly dataset loading rather than training. It's a smoke test that the MEDS path works with the existing stack, not a benchmark. Happy to expose the cap as a task parameter if that'd be generally useful.

Follow-up with a local run on the PhysioNet demo (CPU, ~2.5 min total):

Mortality task samples (train, codes tail-truncated to 256): 238 Epoch 0 / 1: 100%|...| 7/7 [~27s] Test metrics (smoke test only — 1 epoch, tail-truncated sequences; not interpretable as model quality): {'pr_auc': 0.13, 'roc_auc': 0.38, 'f1': 0.22, 'loss': 0.67}

(Metrics are noise on this tiny demo; included only to show the path runs end-to-end.)

Co-authored-by: Cursor <cursoragent@cursor.com>

…rize docstring Co-authored-by: Cursor <cursoragent@cursor.com>

Co-authored-by: Cursor <cursoragent@cursor.com>

AxelNoun · 2026-07-27T13:29:47Z

Thanks for the review / all four addressed and pushed.

…cs output Co-authored-by: Cursor <cursoragent@cursor.com>

jhnwu3

1 last request, you can probably delete it or write a separate utiltiy function that does it for the SampleDataset instead.

jhnwu3 · 2026-07-27T18:11:09Z

+    @staticmethod
+    def summarize(samples: List[Dict[str, Any]]) -> Dict[str, Any]:
+        """Summary statistics of a generated sample set, for cohort reporting.
+


Oh so literally .summarize() before set_task()?

Hmm, I think you might need to change how it works, because it'll be really slow for much larger datasets as it doesn't reuse the set_task() operations and it won't be memory efficient since it's not using our lazyloading?. I think it's easier to do it post-tokenization here, because then you can leverage the parallel threading to do the extraction first and then everything is already a tensor.

Co-authored-by: Cursor <cursoragent@cursor.com>

AxelNoun · 2026-07-27T19:49:31Z

1 last request, you can probably delete it or write a separate utiltiy function that does it for the SampleDataset instead.

I preferred to delete it seemed more coherent to see to add later maybe?

AxelNoun and others added 6 commits July 15, 2026 21:06

docs(examples): add MEDS example and API docs

d8bac4a

Document MEDSDataset in the API reference and add an end-to-end examples/meds_demo.py against the public PhysioNet MIMIC-IV MEDS demo.

fix(tasks): re-export InHospitalMortalityMEDS with explicit alias

8d75571

Satisfies ruff F401 on the newly added __init__ line under tools/check_pr_rules scoped lint (same pattern as eegbci). Co-authored-by: Cursor <cursoragent@cursor.com>

joshuasteier self-requested a review July 17, 2026 20:57

AxelNoun marked this pull request as ready for review July 18, 2026 13:18

jhnwu3 requested changes Jul 26, 2026

View reviewed changes

AxelNoun and others added 3 commits July 27, 2026 15:28

docs: link MEDS schema docs for subject_splits mapping

7bebb45

Co-authored-by: Cursor <cursoragent@cursor.com>

refactor: rename _reconstruct_stays to _group_stays and clarify summa…

070caa1

…rize docstring Co-authored-by: Cursor <cursoragent@cursor.com>

docs: add end-to-end RNN training to MEDS demo

9032e3d

Co-authored-by: Cursor <cursoragent@cursor.com>

docs: link subject_splits in MEDSDataset API page; qualify demo metri…

fea8659

…cs output Co-authored-by: Cursor <cursoragent@cursor.com>

jhnwu3 requested changes Jul 27, 2026

View reviewed changes

refactor: drop summarize helper from InHospitalMortalityMEDS

141ba2f

Co-authored-by: Cursor <cursoragent@cursor.com>

Uh oh!

Conversation

AxelNoun commented Jul 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Changes

Testing

How to verify

Notes for review

Uh oh!

jhnwu3 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

AxelNoun Jul 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

AxelNoun Jul 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

AxelNoun Jul 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

AxelNoun commented Jul 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jhnwu3 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

AxelNoun commented Jul 27, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

AxelNoun commented Jul 15, 2026 •

edited

Loading

AxelNoun Jul 27, 2026 •

edited

Loading

AxelNoun Jul 27, 2026 •

edited

Loading

AxelNoun Jul 27, 2026 •

edited

Loading

AxelNoun commented Jul 27, 2026 •

edited

Loading