Add MEDS dataset support (MEDSDataset + typed Parquet scan path)

[AxelNoun]

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.
• 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: end-to-end on the public demo (download instructions included).

5. InHospitalMortalityMEDS (MEDS-native task)

• One sample per completed stay: reconstruct 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 -v: 13 passed (incl. 4 schema-guard subtests) — Windows, Python 3.12.
• python -m pytest tests/core/test_in_hospital_mortality_meds.py -v: synthetic suite + demo smoke passed (Windows, Python 3.12).
• Full core suite (1,200 tests on this branch / 1,187 on clean master, run via unittest — make is unavailable on Windows): the only failures are environment-specific and unrelated to this change — temp-dir teardown PermissionError on Windows (WinError 32), and network-gated GCS synthetic-MIMIC loads whose URLs are mangled by Windows path joining. By construction the CSV path is unchanged: the datetime fast-path can only trigger on typed (non-CSV) sources.
• The same failures occur on a clean upstream master checkout on the same machine.
• Local smoke on the full public demo: 100 subjects / 803,992 events; canonical dtypes verified (patient_id string, timestamp datetime64[ms]); set_task exercised with a minimal task.
• Cohort via real task pipeline (examples/verify_meds_mortality.py --root <demo>): 12 positive / 238 stays (rate 0.0504); set_task sample count 238.
• Design choices were validated empirically at the pinned versions (dask 2025.11, pyarrow 22, polars 1.35, pandas 2.3): recursive sharded-directory reads, NaT/microsecond behavior of the datetime fast-path, and the silent-parse hazard of date-like integer timestamps (hence the schema guard).

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 + example 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.

Opened as a draft; next updates here: a tutorial notebook.