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

.gitignore

@@ -140,4 +140,18 @@ data/physionet.org/
 .codex
 
 # Model weight files (large binaries, distributed separately)
-weightfiles/
+weightfiles/
+
+# Local / personal (never commit)
+CLAUDE.local.md
+.claude/settings.local.json
+
+# Local test fixtures (download separately; not part of the repo)
+test-resources/core/chestxray14/
+test-resources/meds_demo/
+
+# Local Python environments (not .venv — that pattern is already ignored)
+.venv312/
+
+# Tool caches
+.ruff_cache/

docs/api/datasets.rst

@@ -225,6 +225,7 @@ Available Datasets
     datasets/pyhealth.datasets.MIMIC3Dataset
     datasets/pyhealth.datasets.MIMIC4Dataset
     datasets/pyhealth.datasets.FHIRDataset
+    datasets/pyhealth.datasets.MEDSDataset
     datasets/pyhealth.datasets.MIMIC4FHIR
     datasets/pyhealth.datasets.MedicalTranscriptionsDataset
     datasets/pyhealth.datasets.CardiologyDataset

docs/api/datasets/pyhealth.datasets.MEDSDataset.rst

@@ -0,0 +1,9 @@
+pyhealth.datasets.MEDSDataset
+===================================
+
+Dataset class for data in the `Medical Event Data Standard (MEDS) <https://github.com/Medical-Event-Data-Standard/meds>`_, a minimal event-based schema for machine learning over EHR data (Arnrich et al., ICLR 2024 Workshop on Learning from Time Series For Health). Sharded Parquet event files are read with their native types, and standard MEDS splits (train / tuning / held_out) can be selected directly via the ``subset`` argument.
+
+.. autoclass:: pyhealth.datasets.MEDSDataset
+    :members:
+    :undoc-members:
+    :show-inheritance:

docs/api/tasks.rst

@@ -207,6 +207,7 @@ Available Tasks
 
     Base Task <tasks/pyhealth.tasks.BaseTask>
     In-Hospital Mortality (MIMIC-IV) <tasks/pyhealth.tasks.InHospitalMortalityMIMIC4>
+    In-Hospital Mortality (MEDS) <tasks/pyhealth.tasks.InHospitalMortalityMEDS>
     MIMIC-III ICD-9 Coding <tasks/pyhealth.tasks.MIMIC3ICD9Coding>
     Cardiology Detection <tasks/pyhealth.tasks.cardiology_detect>
     COVID-19 CXR Classification <tasks/pyhealth.tasks.COVID19CXRClassification>

docs/api/tasks/pyhealth.tasks.InHospitalMortalityMEDS.rst

@@ -0,0 +1,7 @@
+pyhealth.tasks.InHospitalMortalityMEDS
+======================================
+
+.. autoclass:: pyhealth.tasks.in_hospital_mortality_meds.InHospitalMortalityMEDS
+    :members:
+    :undoc-members:
+    :show-inheritance:

examples/meds_demo.py

@@ -0,0 +1,76 @@
+"""End-to-end example: loading a MEDS dataset with PyHealth.
+
+This example uses the public *MIMIC-IV demo data in the Medical Event Data
+Standard (MEDS)* (PhysioNet, v0.0.1, ODbL v1.0, ~100 subjects):
+https://doi.org/10.13026/t2y8-ea41
+
+Download it once (open access, ~a few MB):
+
+    wget -r -N -c -np https://physionet.org/files/mimic-iv-demo-meds/0.0.1/
+
+Then run:
+
+    python examples/meds_demo.py \\
+        --root physionet.org/files/mimic-iv-demo-meds/0.0.1
+
+Any dataset following the MEDS layout (``data/**.parquet`` +
+``metadata/subject_splits.parquet``) works the same way. See the MEDS
+specification: https://github.com/Medical-Event-Data-Standard/meds
+"""
+
+import argparse
+
+import polars as pl
+
+from pyhealth.datasets import MEDSDataset
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--root",
+        required=True,
+        help="Root of the MEDS dataset (directory containing data/ and "
+        "metadata/)",
+    )
+    parser.add_argument(
+        "--subset",
+        default="train",
+        help="Split to load as a subset (default: train)",
+    )
+    args = parser.parse_args()
+
+    # 1) Load the full dataset: every Parquet shard under data/ is read,
+    #    including nested split directories (data/<split>/<shard>.parquet).
+    dataset = MEDSDataset(root=args.root)
+    dataset.stats()
+
+    # 2) Peek at the canonical event frame (typed straight from Parquet:
+    #    string patient ids, datetime64[ms] timestamps, float values).
+    events = dataset.global_event_df
+    print(events.head(5).collect())
+
+    # 3) Load a split-restricted subset. Subjects are selected through the
+    #    metadata/subject_splits.parquet assignment; each subset uses its
+    #    own processing cache.
+    subset = MEDSDataset(root=args.root, subset=args.subset)
+    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}")
+
+    # 4) Static (null-time) MEDS events, e.g. demographics, are preserved.
+    n_static = (
+        events.filter(pl.col("timestamp").is_null())
+        .select(pl.len())
+        .collect()
+        .item()
+    )
+    print(f"Static (null-time) events: {n_static}")
+
+    # From here, the dataset behaves like any other PyHealth dataset: use
+    # `dataset.set_task(...)` with an existing task to build samples.
+
+
+if __name__ == "__main__":
+    # BaseDataset spawns Dask worker processes; keep the main-module guard.
+    main()

examples/verify_meds_mortality.py

@@ -0,0 +1,69 @@
+"""Standalone verification of the MEDS in-hospital mortality cohort.
+
+Prints the cohort summary produced by ``InHospitalMortalityMEDS`` on a MEDS
+dataset, so the positive rate (expected ~12/238 on the public MIMIC-IV demo)
+can be confirmed through the actual task pipeline rather than only at the raw
+Parquet level.
+
+Usage:
+    # Download the public demo once (open access, ODbL v1.0):
+    #   wget -r -N -c -np https://physionet.org/files/mimic-iv-demo-meds/0.0.1/
+    python verify_meds_mortality.py \\
+        --root physionet.org/files/mimic-iv-demo-meds/0.0.1
+
+The task needs hadm_id; this script uses the bundled
+``configs/meds_with_hadm.yaml`` automatically.
+"""
+
+import argparse
+from pathlib import Path
+
+import pyhealth.datasets.configs as meds_configs
+from pyhealth.datasets import MEDSDataset
+from pyhealth.tasks import InHospitalMortalityMEDS
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--root",
+        required=True,
+        help="Root of the MEDS dataset (contains data/ and metadata/).",
+    )
+    parser.add_argument(
+        "--observation-window",
+        default="full_stay",
+        choices=["full_stay", "first_hours"],
+    )
+    parser.add_argument("--window-hours", type=float, default=48.0)
+    args = parser.parse_args()
+
+    cfg = Path(meds_configs.__file__).parent / "meds_with_hadm.yaml"
+    dataset = MEDSDataset(root=args.root, config_path=str(cfg))
+    task = InHospitalMortalityMEDS(
+        observation_window=args.observation_window,
+        window_hours=args.window_hours,
+    )
+
+    # Apply per patient to keep raw (untokenized) samples, so the label
+    # counts are directly inspectable. set_task would tokenize the codes.
+    samples = []
+    for patient_id in dataset.unique_patient_ids:
+        samples.extend(task(dataset.get_patient(patient_id)))
+
+    summary = InHospitalMortalityMEDS.summarize(samples)
+    print(f"root                 : {args.root}")
+    print(f"observation_window   : {args.observation_window}")
+    print(f"n_samples (stays)    : {summary['n_samples']}")
+    print(f"n_patients           : {summary['n_patients']}")
+    print(f"n_positive (died)    : {summary['n_positive']}")
+    print(f"positive_rate        : {summary['positive_rate']:.4f}")
+    print(f"mean_sequence_length : {summary['mean_sequence_length']:.1f}")
+
+    # Cross-check the intended user path returns the same sample count.
+    sample_dataset = dataset.set_task(task)
+    print(f"set_task sample count : {len(sample_dataset)} (should equal n_samples)")
+
+
+if __name__ == "__main__":
+    main()

pyhealth/datasets/__init__.py

@@ -57,6 +57,7 @@ def __init__(self, *args, **kwargs):
 from .eicu import eICUDataset
 from .isruc import ISRUCDataset
 from .medical_transcriptions import MedicalTranscriptionsDataset
+from .meds import MEDSDataset as MEDSDataset  # noqa: E402
 from .mimic3 import MIMIC3Dataset
 from .mimic4 import MIMIC4CXRDataset, MIMIC4Dataset, MIMIC4EHRDataset, MIMIC4NoteDataset
 from .fhir import FHIRDataset, MIMIC4FHIR

pyhealth/datasets/base_dataset.py

@@ -315,6 +315,15 @@ class BaseDataset(ABC):
         config (dict): Configuration loaded from a YAML file.
         global_event_df (pl.LazyFrame): The global event data frame.
         dev (bool): Whether to enable dev mode (limit to 1000 patients).
+
+    Examples:
+        >>> from pyhealth.datasets import BaseDataset
+        >>> dataset = BaseDataset(
+        ...     root="/path/to/source",
+        ...     tables=["patients", "diagnoses"],
+        ...     config_path="/path/to/config.yaml",
+        ... )
+        >>> dataset.stats()
     """
 
     def __init__(
@@ -425,6 +434,68 @@ def clean_tmpdir(self) -> None:
         if tmp_dir.exists():
             shutil.rmtree(tmp_dir)
 
+    def _scan_table(self, source_path: str) -> dd.DataFrame:
+        """Routes a table source to the appropriate scanner based on its format.
+
+        Parquet sources (``.parquet``/``.pq`` files, glob patterns targeting
+        such files, or directories of Parquet shards) are handled by
+        :meth:`_scan_parquet`. Any other source falls back to the existing
+        CSV/TSV(.gz) scanner, preserving prior behavior for all datasets.
+
+        Args:
+            source_path (str): Path to the table source.
+
+        Returns:
+            dd.DataFrame: The Dask DataFrame for the table source.
+        """
+        stripped = source_path.rstrip("/")
+        if stripped.endswith((".parquet", ".pq")) or (
+            not is_url(source_path) and Path(source_path).is_dir()
+        ):
+            return self._scan_parquet(source_path)
+        return self._scan_csv_tsv_gz(source_path)
+
+    def _scan_parquet(self, source_path: str) -> dd.DataFrame:
+        """Scans a Parquet source and returns a Dask DataFrame.
+
+        The source may be a single ``.parquet``/``.pq`` file, a glob pattern,
+        or a directory that is scanned recursively — which supports sharded
+        datasets such as MEDS, laid out as ``data/<split>/<shard>.parquet``.
+
+        Unlike :meth:`_scan_csv_tsv_gz`, no all-string schema coercion is
+        applied: Parquet files embed their schema, so source dtypes (native
+        timestamps, numeric columns, nullable strings) are preserved and
+        handled downstream by :meth:`load_table`.
+
+        Args:
+            source_path (str): Path to a Parquet file, directory, or glob.
+
+        Returns:
+            dd.DataFrame: The Dask DataFrame backed by the Parquet source.
+
+        Raises:
+            FileNotFoundError: If the source path does not exist, or if a
+                directory source contains no Parquet files.
+        """
+        path = Path(source_path)
+        is_glob = any(ch in source_path for ch in "*?[")
+        if not is_glob:
+            if not path.exists():
+                raise FileNotFoundError(
+                    f"Parquet source does not exist: {source_path}"
+                )
+            if path.is_dir() and not any(
+                itertools.chain(path.rglob("*.parquet"), path.rglob("*.pq"))
+            ):
+                raise FileNotFoundError(
+                    f"Directory contains no Parquet files: {source_path}"
+                )
+        return dd.read_parquet(
+            source_path,
+            split_row_groups=True,  # type: ignore
+            blocksize="64MB",
+        )
+
     def _scan_csv_tsv_gz(self, source_path: str) -> dd.DataFrame:
         """Scans a CSV/TSV file (possibly gzipped) and returns a Dask DataFrame.
 
@@ -596,7 +667,8 @@ def load_table(self, table_name: str) -> dd.DataFrame:
 
         Raises:
             ValueError: If the table is not found in the config.
-            FileNotFoundError: If the CSV file for the table or join is not found.
+            FileNotFoundError: If the source file (CSV/TSV or Parquet) for the
+                table or join is not found.
         """
         assert self.config is not None, "Config must be provided to load tables"
 
@@ -608,7 +680,7 @@ def load_table(self, table_name: str) -> dd.DataFrame:
         csv_path = clean_path(csv_path)
 
         logger.info(f"Scanning table: {table_name} from {csv_path}")
-        df = self._scan_csv_tsv_gz(csv_path)
+        df = self._scan_table(csv_path)
 
         # Convert column names to lowercase before calling preprocess_func
         df = df.rename(columns=str.lower)
@@ -627,7 +699,7 @@ def load_table(self, table_name: str) -> dd.DataFrame:
             other_csv_path = f"{self.root}/{join_cfg.file_path}"
             other_csv_path = clean_path(other_csv_path)
             logger.info(f"Joining with table: {other_csv_path}")
-            join_df = self._scan_csv_tsv_gz(other_csv_path)
+            join_df = self._scan_table(other_csv_path)
             join_df = join_df.rename(columns=str.lower)
             join_key = join_cfg.on
             columns = join_cfg.columns
@@ -651,14 +723,21 @@ def load_table(self, table_name: str) -> dd.DataFrame:
                 timestamp_series: dd.Series = functools.reduce(
                     operator.add, (df[col].astype("string") for col in timestamp_col)
                 )
+                timestamp_series = dd.to_datetime(
+                    timestamp_series,
+                    format=timestamp_format,
+                    errors="raise",
+                )
+            elif pd.api.types.is_datetime64_any_dtype(df[timestamp_col].dtype):
+                # Typed sources (e.g. Parquet) already carry native timestamps:
+                # skip the string round-trip and only normalize the unit below.
+                timestamp_series: dd.Series = df[timestamp_col]
             else:
-                timestamp_series: dd.Series = df[timestamp_col].astype("string")
-
-            timestamp_series: dd.Series = dd.to_datetime(
-                timestamp_series,
-                format=timestamp_format,
-                errors="raise",
-            )
+                timestamp_series = dd.to_datetime(
+                    df[timestamp_col].astype("string"),
+                    format=timestamp_format,
+                    errors="raise",
+                )
             df: dd.DataFrame = df.assign(
                 timestamp=timestamp_series.astype("datetime64[ms]")
             )

pyhealth/datasets/configs/meds.yaml

@@ -0,0 +1,34 @@
+# MEDS (Medical Event Data Standard) tables.
+#
+# MEDS is already flat (one row per measurement), so a single event table
+# covers the data shards. The canonical subject-to-split mapping is exposed
+# as a second, ordinary event table -- the same pattern as the `splits`
+# table in ehrshot.yaml. TableConfig has no format field: Parquet reading is
+# handled by BaseDataset._scan_table / _scan_parquet, so no config-schema
+# change is needed for the 21 existing datasets.
+#
+# Paths match mimic-iv-demo-meds: data/<split>/*.parquet +
+# metadata/subject_splits.parquet.
+version: "1.0"
+tables:
+  meds:
+    # Directory of shards; dd.read_parquet reads the nested tree whole.
+    file_path: "data"
+    patient_id: "subject_id"
+    # Native datetime64[us] in MEDS Parquet, narrowed to [ms] by the
+    # BaseDataset typed-timestamp fast-path. No timestamp_format: that
+    # would imply text parsing. MEDSDataset rejects non-timestamp /
+    # tz-aware columns at construction (footer schema guard).
+    timestamp: "time"
+    attributes:
+      - "code"
+      - "numeric_value"
+
+  # Canonical split map as events (attribute `subject_splits/split`).
+  # Opt-in: tables=["meds", "subject_splits"].
+  subject_splits:
+    file_path: "metadata/subject_splits.parquet"
+    patient_id: "subject_id"
+    timestamp: null
+    attributes:
+      - "split"