Repository: mmcdermott/MEDS_transforms

Issue:
Extraction ETL crashes if you include the `extract_metadata` stage but you don't have any `_metadata` blocks in your configs.
Instead, this stage should just copy or symlink over any existing `metadata/codes.parquet` and terminate in this case.

**In case this issue is impacting anybody**, before it gets formally fixed you can solve this by just removing this stage from your extraction pipeline. You can do this on the command line on the fly by just overwriting the `stages` parameter in the normal Hydra manner to include all the other stages, but not this stage (and then also skip the stage-specific script as well, naturally). See https://github.com/mmcdermott/MEDS_transforms?tab=readme-ov-file#notes for an example of this overwrite syntax.

Discussion:

Patch:
diff --git a/src/MEDS_transforms/extract/extract_code_metadata.py b/src/MEDS_transforms/extract/extract_code_metadata.py
index 1b8b394..e9133eb 100644
--- a/src/MEDS_transforms/extract/extract_code_metadata.py
+++ b/src/MEDS_transforms/extract/extract_code_metadata.py
@@ -364,6 +364,10 @@ def main(cfg: DictConfig):
     OmegaConf.save(event_conversion_cfg, partial_metadata_dir / "event_conversion_config.yaml")
 
     events_and_metadata_by_metadata_fp = get_events_and_metadata_by_metadata_fp(event_conversion_cfg)
+    if not events_and_metadata_by_metadata_fp:
+        logger.info("No _metadata blocks in the event_conversion_config.yaml found. Exiting...")
+        return
+
     event_metadata_configs = list(events_and_metadata_by_metadata_fp.items())
     random.shuffle(event_metadata_configs)
 


Test patch:
diff --git a/tests/test_extract_no_metadata.py b/tests/test_extract_no_metadata.py
new file mode 100644
index 0000000..f1945af
--- /dev/null
+++ b/tests/test_extract_no_metadata.py
@@ -0,0 +1,631 @@
+"""Tests the full end-to-end extraction process.
+
+Set the bash env variable `DO_USE_LOCAL_SCRIPTS=1` to use the local py files, rather than the installed
+scripts.
+"""
+
+import os
+
+import rootutils
+
+root = rootutils.setup_root(__file__, dotenv=True, pythonpath=True, cwd=True)
+
+code_root = root / "src" / "MEDS_transforms"
+extraction_root = code_root / "extract"
+
+if os.environ.get("DO_USE_LOCAL_SCRIPTS", "0") == "1":
+    SHARD_EVENTS_SCRIPT = extraction_root / "shard_events.py"
+    SPLIT_AND_SHARD_SCRIPT = extraction_root / "split_and_shard_patients.py"
+    CONVERT_TO_SHARDED_EVENTS_SCRIPT = extraction_root / "convert_to_sharded_events.py"
+    MERGE_TO_MEDS_COHORT_SCRIPT = extraction_root / "merge_to_MEDS_cohort.py"
+    EXTRACT_CODE_METADATA_SCRIPT = extraction_root / "extract_code_metadata.py"
+    FINALIZE_DATA_SCRIPT = extraction_root / "finalize_MEDS_data.py"
+    FINALIZE_METADATA_SCRIPT = extraction_root / "finalize_MEDS_metadata.py"
+else:
+    SHARD_EVENTS_SCRIPT = "MEDS_extract-shard_events"
+    SPLIT_AND_SHARD_SCRIPT = "MEDS_extract-split_and_shard_patients"
+    CONVERT_TO_SHARDED_EVENTS_SCRIPT = "MEDS_extract-convert_to_sharded_events"
+    MERGE_TO_MEDS_COHORT_SCRIPT = "MEDS_extract-merge_to_MEDS_cohort"
+    EXTRACT_CODE_METADATA_SCRIPT = "MEDS_extract-extract_code_metadata"
+    FINALIZE_DATA_SCRIPT = "MEDS_extract-finalize_MEDS_data"
+    FINALIZE_METADATA_SCRIPT = "MEDS_extract-finalize_MEDS_metadata"
+
+import json
+import tempfile
+from io import StringIO
+from pathlib import Path
+
+import polars as pl
+from meds import __version__ as MEDS_VERSION
+
+from .utils import assert_df_equal, run_command
+
+# Test data (inputs)
+
+SUBJECTS_CSV = """
+MRN,dob,eye_color,height
+1195293,06/20/1978,BLUE,164.6868838269085
+239684,12/28/1980,BROWN,175.271115221764
+1500733,07/20/1986,BROWN,158.60131573580904
+814703,03/28/1976,HAZEL,156.48559093209357
+754281,12/19/1988,BROWN,166.22261567137025
+68729,03/09/1978,HAZEL,160.3953106166676
+"""
+
+ADMIT_VITALS_CSV = """
+patient_id,admit_date,disch_date,department,vitals_date,HR,temp
+239684,"05/11/2010, 17:41:51","05/11/2010, 19:27:19",CARDIAC,"05/11/2010, 18:57:18",112.6,95.5
+754281,"01/03/2010, 06:27:59","01/03/2010, 08:22:13",PULMONARY,"01/03/2010, 06:27:59",142.0,99.8
+814703,"02/05/2010, 05:55:39","02/05/2010, 07:02:30",ORTHOPEDIC,"02/05/2010, 05:55:39",170.2,100.1
+239684,"05/11/2010, 17:41:51","05/11/2010, 19:27:19",CARDIAC,"05/11/2010, 18:25:35",113.4,95.8
+68729,"05/26/2010, 02:30:56","05/26/2010, 04:51:52",PULMONARY,"05/26/2010, 02:30:56",86.0,97.8
+1195293,"06/20/2010, 19:23:52","06/20/2010, 20:50:04",CARDIAC,"06/20/2010, 20:12:31",112.5,99.8
+1500733,"06/03/2010, 14:54:38","06/03/2010, 16:44:26",ORTHOPEDIC,"06/03/2010, 16:20:49",90.1,100.1
+239684,"05/11/2010, 17:41:51","05/11/2010, 19:27:19",CARDIAC,"05/11/2010, 17:48:48",105.1,96.2
+239684,"05/11/2010, 17:41:51","05/11/2010, 19:27:19",CARDIAC,"05/11/2010, 17:41:51",102.6,96.0
+1195293,"06/20/2010, 19:23:52","06/20/2010, 20:50:04",CARDIAC,"06/20/2010, 19:25:32",114.1,100.0
+1500733,"06/03/2010, 14:54:38","06/03/2010, 16:44:26",ORTHOPEDIC,"06/03/2010, 14:54:38",91.4,100.0
+1195293,"06/20/2010, 19:23:52","06/20/2010, 20:50:04",CARDIAC,"06/20/2010, 20:41:33",107.5,100.4
+1195293,"06/20/2010, 19:23:52","06/20/2010, 20:50:04",CARDIAC,"06/20/2010, 20:24:44",107.7,100.0
+1195293,"06/20/2010, 19:23:52","06/20/2010, 20:50:04",CARDIAC,"06/20/2010, 19:45:19",119.8,99.9
+1195293,"06/20/2010, 19:23:52","06/20/2010, 20:50:04",CARDIAC,"06/20/2010, 19:23:52",109.0,100.0
+1500733,"06/03/2010, 14:54:38","06/03/2010, 16:44:26",ORTHOPEDIC,"06/03/2010, 15:39:49",84.4,100.3
+"""
+
+INPUT_METADATA_FILE = """
+lab_code,title,loinc
+HR,Heart Rate,8867-4
+temp,Body Temperature,8310-5
+"""
+
+DEMO_METADATA_FILE = """
+eye_color,description
+BROWN,"Brown Eyes. The most common eye color."
+BLUE,"Blue Eyes. Less common than brown."
+HAZEL,"Hazel eyes. These are uncommon"
+GREEN,"Green eyes. These are rare."
+"""
+
+EVENT_CFGS_YAML = """
+subjects:
+  patient_id_col: MRN
+  eye_color:
+    code:
+      - EYE_COLOR
+      - col(eye_color)
+    time: null
+  height:
+    code: HEIGHT
+    time: null
+    numeric_value: height
+  dob:
+    code: DOB
+    time: col(dob)
+    time_format: "%m/%d/%Y"
+admit_vitals:
+  admissions:
+    code:
+      - ADMISSION
+      - col(department)
+    time: col(admit_date)
+    time_format: "%m/%d/%Y, %H:%M:%S"
+  discharge:
+    code: DISCHARGE
+    time: col(disch_date)
+    time_format: "%m/%d/%Y, %H:%M:%S"
+  HR:
+    code: HR
+    time: col(vitals_date)
+    time_format: "%m/%d/%Y, %H:%M:%S"
+    numeric_value: HR
+  temp:
+    code: TEMP
+    time: col(vitals_date)
+    time_format: "%m/%d/%Y, %H:%M:%S"
+    numeric_value: temp
+"""
+
+# Test data (expected outputs) -- ALL OF THIS MAY CHANGE IF THE SEED OR DATA CHANGES
+EXPECTED_SPLITS = {
+    "train/0": [239684, 1195293],
+    "train/1": [68729, 814703],
+    "tuning/0": [754281],
+    "held_out/0": [1500733],
+}
+
+PATIENT_SPLITS_DF = pl.DataFrame(
+    {
+        "patient_id": [239684, 1195293, 68729, 814703, 754281, 1500733],
+        "split": ["train", "train", "train", "train", "tuning", "held_out"],
+    }
+)
+
+
+def get_expected_output(df: str) -> pl.DataFrame:
+    return (
+        pl.read_csv(source=StringIO(df))
+        .select(
+            "patient_id",
+            pl.col("time").str.strptime(pl.Datetime, "%m/%d/%Y, %H:%M:%S").alias("time"),
+            pl.col("code"),
+            "numeric_value",
+        )
+        .sort(by=["patient_id", "time"])
+    )
+
+
+MEDS_OUTPUT_TRAIN_0_SUBJECTS = """
+patient_id,time,code,numeric_value
+239684,,EYE_COLOR//BROWN,
+239684,,HEIGHT,175.271115221764
+239684,"12/28/1980, 00:00:00",DOB,
+1195293,,EYE_COLOR//BLUE,
+1195293,,HEIGHT,164.6868838269085
+1195293,"06/20/1978, 00:00:00",DOB,
+"""
+
+MEDS_OUTPUT_TRAIN_0_ADMIT_VITALS = """
+patient_id,time,code,numeric_value
+239684,"05/11/2010, 17:41:51",ADMISSION//CARDIAC,
+239684,"05/11/2010, 17:41:51",HR,102.6
+239684,"05/11/2010, 17:41:51",TEMP,96.0
+239684,"05/11/2010, 17:48:48",HR,105.1
+239684,"05/11/2010, 17:48:48",TEMP,96.2
+239684,"05/11/2010, 18:25:35",HR,113.4
+239684,"05/11/2010, 18:25:35",TEMP,95.8
+239684,"05/11/2010, 18:57:18",HR,112.6
+239684,"05/11/2010, 18:57:18",TEMP,95.5
+239684,"05/11/2010, 19:27:19",DISCHARGE,
+1195293,"06/20/2010, 19:23:52",ADMISSION//CARDIAC,
+1195293,"06/20/2010, 19:23:52",HR,109.0
+1195293,"06/20/2010, 19:23:52",TEMP,100.0
+1195293,"06/20/2010, 19:25:32",HR,114.1
+1195293,"06/20/2010, 19:25:32",TEMP,100.0
+1195293,"06/20/2010, 19:45:19",HR,119.8
+1195293,"06/20/2010, 19:45:19",TEMP,99.9
+1195293,"06/20/2010, 20:12:31",HR,112.5
+1195293,"06/20/2010, 20:12:31",TEMP,99.8
+1195293,"06/20/2010, 20:24:44",HR,107.7
+1195293,"06/20/2010, 20:24:44",TEMP,100.0
+1195293,"06/20/2010, 20:41:33",HR,107.5
+1195293,"06/20/2010, 20:41:33",TEMP,100.4
+1195293,"06/20/2010, 20:50:04",DISCHARGE,
+"""
+
+MEDS_OUTPUT_TRAIN_1_SUBJECTS = """
+patient_id,time,code,numeric_value
+68729,,EYE_COLOR//HAZEL,
+68729,,HEIGHT,160.3953106166676
+68729,"03/09/1978, 00:00:00",DOB,
+814703,,EYE_COLOR//HAZEL,
+814703,,HEIGHT,156.48559093209357
+814703,"03/28/1976, 00:00:00",DOB,
+"""
+
+MEDS_OUTPUT_TRAIN_1_ADMIT_VITALS = """
+patient_id,time,code,numeric_value
+68729,"05/26/2010, 02:30:56",ADMISSION//PULMONARY,
+68729,"05/26/2010, 02:30:56",HR,86.0
+68729,"05/26/2010, 02:30:56",TEMP,97.8
+68729,"05/26/2010, 04:51:52",DISCHARGE,
+814703,"02/05/2010, 05:55:39",ADMISSION//ORTHOPEDIC,
+814703,"02/05/2010, 05:55:39",HR,170.2
+814703,"02/05/2010, 05:55:39",TEMP,100.1
+814703,"02/05/2010, 07:02:30",DISCHARGE,
+"""
+
+MEDS_OUTPUT_TUNING_0_SUBJECTS = """
+patient_id,time,code,numeric_value
+754281,,EYE_COLOR//BROWN,
+754281,,HEIGHT,166.22261567137025
+754281,"12/19/1988, 00:00:00",DOB,
+"""
+
+MEDS_OUTPUT_TUNING_0_ADMIT_VITALS = """
+patient_id,time,code,numeric_value
+754281,"01/03/2010, 06:27:59",ADMISSION//PULMONARY,
+754281,"01/03/2010, 06:27:59",HR,142.0
+754281,"01/03/2010, 06:27:59",TEMP,99.8
+754281,"01/03/2010, 08:22:13",DISCHARGE,
+"""
+
+MEDS_OUTPUT_HELD_OUT_0_SUBJECTS = """
+patient_id,time,code,numeric_value
+1500733,,EYE_COLOR//BROWN,
+1500733,,HEIGHT,158.60131573580904
+1500733,"07/20/1986, 00:00:00",DOB,
+"""
+
+MEDS_OUTPUT_HELD_OUT_0_ADMIT_VITALS = """
+patient_id,time,code,numeric_value
+1500733,"06/03/2010, 14:54:38",ADMISSION//ORTHOPEDIC,
+1500733,"06/03/2010, 14:54:38",HR,91.4
+1500733,"06/03/2010, 14:54:38",TEMP,100.0
+1500733,"06/03/2010, 15:39:49",HR,84.4
+1500733,"06/03/2010, 15:39:49",TEMP,100.3
+1500733,"06/03/2010, 16:20:49",HR,90.1
+1500733,"06/03/2010, 16:20:49",TEMP,100.1
+1500733,"06/03/2010, 16:44:26",DISCHARGE,
+"""
+
+MEDS_OUTPUT_CODE_METADATA_FILE = """
+code,description,parent_codes
+"""
+
+MEDS_OUTPUT_DATASET_METADATA_JSON = {
+    "dataset_name": "TEST",
+    "dataset_version": "1.0",
+    "etl_name": "MEDS_transforms",
+    # "etl_version": None,  # We don't test this as it changes with the commits.
+    "meds_version": MEDS_VERSION,
+}
+
+SUB_SHARDED_OUTPUTS = {
+    "train/0": {
+        "subjects": MEDS_OUTPUT_TRAIN_0_SUBJECTS,
+        "admit_vitals": MEDS_OUTPUT_TRAIN_0_ADMIT_VITALS,
+    },
+    "train/1": {
+        "subjects": MEDS_OUTPUT_TRAIN_1_SUBJECTS,
+        "admit_vitals": MEDS_OUTPUT_TRAIN_1_ADMIT_VITALS,
+    },
+    "tuning/0": {
+        "subjects": MEDS_OUTPUT_TUNING_0_SUBJECTS,
+        "admit_vitals": MEDS_OUTPUT_TUNING_0_ADMIT_VITALS,
+    },
+    "held_out/0": {
+        "subjects": MEDS_OUTPUT_HELD_OUT_0_SUBJECTS,
+        "admit_vitals": MEDS_OUTPUT_HELD_OUT_0_ADMIT_VITALS,
+    },
+}
+
+
+MEDS_OUTPUTS = {
+    "train/0": [MEDS_OUTPUT_TRAIN_0_SUBJECTS, MEDS_OUTPUT_TRAIN_0_ADMIT_VITALS],
+    "train/1": [MEDS_OUTPUT_TRAIN_1_SUBJECTS, MEDS_OUTPUT_TRAIN_1_ADMIT_VITALS],
+    "tuning/0": [MEDS_OUTPUT_TUNING_0_SUBJECTS, MEDS_OUTPUT_TUNING_0_ADMIT_VITALS],
+    "held_out/0": [MEDS_OUTPUT_HELD_OUT_0_SUBJECTS, MEDS_OUTPUT_HELD_OUT_0_ADMIT_VITALS],
+}
+
+
+def test_extraction():
+    with tempfile.TemporaryDirectory() as d:
+        raw_cohort_dir = Path(d) / "raw_cohort"
+        MEDS_cohort_dir = Path(d) / "MEDS_cohort"
+
+        # Create the directories
+        raw_cohort_dir.mkdir()
+        MEDS_cohort_dir.mkdir()
+
+        subjects_csv = raw_cohort_dir / "subjects.csv"
+        admit_vitals_csv = raw_cohort_dir / "admit_vitals.csv"
+        event_cfgs_yaml = raw_cohort_dir / "event_cfgs.yaml"
+
+        demo_metadata_csv = raw_cohort_dir / "demo_metadata.csv"
+        input_metadata_csv = raw_cohort_dir / "input_metadata.csv"
+
+        # Write the CSV files
+        subjects_csv.write_text(SUBJECTS_CSV.strip())
+        admit_vitals_csv.write_text(ADMIT_VITALS_CSV.strip())
+        demo_metadata_csv.write_text(DEMO_METADATA_FILE.strip())
+        input_metadata_csv.write_text(INPUT_METADATA_FILE.strip())
+
+        # Mix things up -- have one CSV be also in parquet format.
+        admit_vitals_parquet = raw_cohort_dir / "admit_vitals.parquet"
+        df = pl.read_csv(admit_vitals_csv)
+
+        df.write_parquet(admit_vitals_parquet, use_pyarrow=True)
+
+        # Write the event config YAML
+        event_cfgs_yaml.write_text(EVENT_CFGS_YAML)
+
+        # Run the extraction script
+        #   1. Sub-shard the data (this will be a null operation in this case, but it is worth doing just in
+        #      case.
+        #   2. Collect the patient splits.
+        #   3. Extract the events and sub-shard by patient.
+        #   4. Merge to the final output.
+
+        extraction_config_kwargs = {
+            "input_dir": str(raw_cohort_dir.resolve()),
+            "cohort_dir": str(MEDS_cohort_dir.resolve()),
+            "event_conversion_config_fp": str(event_cfgs_yaml.resolve()),
+            "stage_configs.split_and_shard_patients.split_fracs.train": 4 / 6,
+            "stage_configs.split_and_shard_patients.split_fracs.tuning": 1 / 6,
+            "stage_configs.split_and_shard_patients.split_fracs.held_out": 1 / 6,
+            "stage_configs.shard_events.row_chunksize": 10,
+            "stage_configs.split_and_shard_patients.n_patients_per_shard": 2,
+            "hydra.verbose": True,
+            "etl_metadata.dataset_name": "TEST",
+            "etl_metadata.dataset_version": "1.0",
+        }
+
+        all_stderrs = []
+        all_stdouts = []
+
+        # Stage 1: Sub-shard the data
+        stderr, stdout = run_command(SHARD_EVENTS_SCRIPT, extraction_config_kwargs, "shard_events")
+
+        all_stderrs.append(stderr)
+        all_stdouts.append(stdout)
+
+        subsharded_dir = MEDS_cohort_dir / "shard_events"
+
+        try:
+            out_files = list(subsharded_dir.glob("**/*.parquet"))
+            assert len(out_files) == 3, f"Expected 3 output files, got {len(out_files)}."
+
+            # Checking specific out files:
+            #   1. subjects.parquet
+            subjects_out = subsharded_dir / "subjects" / "[0-6).parquet"
+            assert subjects_out.is_file(), f"Expected {subjects_out} to exist. Files include {out_files}."
+
+            assert_df_equal(
+                pl.read_csv(subjects_csv),
+                pl.read_parquet(subjects_out, glob=False),
+                "Subjects should be equal after sub-sharding",
+                check_column_order=False,
+                check_row_order=False,
+            )
+        except AssertionError as e:
+            full_stderr = "\n".join(all_stderrs)
+            print("Sub-sharding failed")
+            print(f"stderr:\n{full_stderr}")
+            raise e
+
+        #   2. admit_vitals.parquet
+        df_chunks = []
+        for chunk in ["[0-10)", "[10-16)"]:
+            admit_vitals_chunk_fp = subsharded_dir / "admit_vitals" / f"{chunk}.parquet"
+            assert admit_vitals_chunk_fp.is_file(), f"Expected {admit_vitals_chunk_fp} to exist."
+
+            df_chunks.append(pl.read_parquet(admit_vitals_chunk_fp, glob=False))
+
+        assert_df_equal(
+            pl.read_csv(admit_vitals_csv),
+            pl.concat(df_chunks),
+            "Admit vitals should be equal after sub-sharding",
+            check_column_order=False,
+            check_row_order=False,
+        )
+
+        # Stage 2: Collect the patient splits
+        stderr, stdout = run_command(
+            SPLIT_AND_SHARD_SCRIPT,
+            extraction_config_kwargs,
+            "split_and_shard_patients",
+        )
+
+        all_stderrs.append(stderr)
+        all_stdouts.append(stdout)
+
+        try:
+            shards_fp = MEDS_cohort_dir / "metadata" / ".shards.json"
+            assert shards_fp.is_file(), f"Expected splits @ {str(shards_fp.resolve())} to exist."
+
+            splits = json.loads(shards_fp.read_text())
+            expected_keys = ["train/0", "train/1", "tuning/0", "held_out/0"]
+
+            expected_keys_str = ", ".join(f"'{k}'" for k in expected_keys)
+            got_keys_str = ", ".join(f"'{k}'" for k in splits.keys())
+
+            assert set(splits.keys()) == set(expected_keys), (
+                f"Expected splits to have keys {expected_keys_str}.\n" f"Got keys: {got_keys_str}"
+            )
+
+            assert splits == EXPECTED_SPLITS, (
+                f"Expected splits to be {EXPECTED_SPLITS}, got {splits}. NOTE THIS MAY CHANGE IF THE SEED OR "
+                "DATA CHANGES -- FAILURE HERE MAY BE JUST DUE TO A NON-DETERMINISTIC SPLIT AND THE TEST "
+                "NEEDING TO BE UPDATED."
+            )
+        except AssertionError as e:
+            print("Failed to split patients")
+            print(f"stderr:\n{stderr}")
+            print(f"stdout:\n{stdout}")
+            raise e
+
+        # Stage 3: Extract the events and sub-shard by patient
+        stderr, stdout = run_command(
+            CONVERT_TO_SHARDED_EVENTS_SCRIPT,
+            extraction_config_kwargs,
+            "convert_events",
+        )
+        all_stderrs.append(stderr)
+        all_stdouts.append(stdout)
+
+        patient_subsharded_folder = MEDS_cohort_dir / "convert_to_sharded_events"
+        assert patient_subsharded_folder.is_dir(), f"Expected {patient_subsharded_folder} to be a directory."
+
+        for split, expected_outputs in SUB_SHARDED_OUTPUTS.items():
+            for prefix, expected_df_L in expected_outputs.items():
+                if not isinstance(expected_df_L, list):
+                    expected_df_L = [expected_df_L]
+
+                expected_df = pl.concat([get_expected_output(df) for df in expected_df_L])
+
+                fps = list((patient_subsharded_folder / split / prefix).glob("*.parquet"))
+                assert len(fps) > 0
+
+                # We add a "unique" here as there may be some duplicates across the row-group sub-shards.
+                got_df = pl.concat([pl.read_parquet(fp, glob=False) for fp in fps]).unique()
+                try:
+                    assert_df_equal(
+                        expected_df,
+                        got_df,
+                        f"Expected output for split {split}/{prefix} to be equal to the expected output.",
+                        check_column_order=False,
+                        check_row_order=False,
+                    )
+                except AssertionError as e:
+                    print(f"Failed on split {split}/{prefix}")
+                    print(f"stderr:\n{stderr}")
+                    print(f"stdout:\n{stdout}")
+                    raise e
+
+        # Stage 4: Merge to the final output
+        stderr, stdout = run_command(
+            MERGE_TO_MEDS_COHORT_SCRIPT,
+            extraction_config_kwargs,
+            "merge_to_MEDS_cohort",
+        )
+        all_stderrs.append(stderr)
+        all_stdouts.append(stdout)
+
+        full_stderr = "\n".join(all_stderrs)
+        full_stdout = "\n".join(all_stdouts)
+
+        # Check the final output
+        output_folder = MEDS_cohort_dir / "merge_to_MEDS_cohort"
+        try:
+            for split, expected_df_L in MEDS_OUTPUTS.items():
+                if not isinstance(expected_df_L, list):
+                    expected_df_L = [expected_df_L]
+
+                expected_df = pl.concat([get_expected_output(df) for df in expected_df_L])
+
+                fp = output_folder / f"{split}.parquet"
+                assert fp.is_file(), f"Expected {fp} to exist.\nstderr:\n{stderr}\nstdout:\n{stdout}"
+
+                got_df = pl.read_parquet(fp, glob=False)
+                assert_df_equal(
+                    expected_df,
+                    got_df,
+                    f"Expected output for split {split} to be equal to the expected output.",
+                    check_column_order=False,
+                    check_row_order=False,
+                )
+
+                assert got_df["patient_id"].is_sorted(), f"Patient IDs should be sorted for split {split}."
+                for subj in splits[split]:
+                    got_df_subj = got_df.filter(pl.col("patient_id") == subj)
+                    assert got_df_subj[
+                        "time"
+                    ].is_sorted(), f"Times should be sorted for patient {subj} in split {split}."
+
+        except AssertionError as e:
+            print(f"Failed on split {split}")
+            print(f"stderr:\n{full_stderr}")
+            print(f"stdout:\n{full_stdout}")
+            raise e
+
+        # Stage 6: Extract code metadata
+        stderr, stdout = run_command(
+            EXTRACT_CODE_METADATA_SCRIPT,
+            extraction_config_kwargs,
+            "extract_code_metadata",
+        )
+        all_stderrs.append(stderr)
+        all_stdouts.append(stdout)
+
+        full_stderr = "\n".join(all_stderrs)
+        full_stdout = "\n".join(all_stdouts)
+
+        output_file = MEDS_cohort_dir / "extract_code_metadata" / "codes.parquet"
+        assert (
+            not output_file.is_file()
+        ), f"Expected {output_file} to not  exist: stderr:\n{stderr}\nstdout:\n{stdout}"
+
+        # Stage 7: Finalize the MEDS data
+        stderr, stdout = run_command(
+            FINALIZE_DATA_SCRIPT,
+            extraction_config_kwargs,
+            "finalize_MEDS_data",
+        )
+        all_stderrs.append(stderr)
+        all_stdouts.append(stdout)
+
+        full_stderr = "\n".join(all_stderrs)
+        full_stdout = "\n".join(all_stdouts)
+
+        # Check the final output
+        output_folder = MEDS_cohort_dir / "data"
+        try:
+            for split, expected_df_L in MEDS_OUTPUTS.items():
+                if not isinstance(expected_df_L, list):
+                    expected_df_L = [expected_df_L]
+
+                expected_df = pl.concat([get_expected_output(df) for df in expected_df_L]).with_columns(
+                    pl.col("numeric_value").cast(pl.Float32)
+                )
+
+                fp = output_folder / f"{split}.parquet"
+                assert fp.is_file(), f"Expected {fp} to exist."
+
+                got_df = pl.read_parquet(fp, glob=False)
+                assert_df_equal(
+                    expected_df,
+                    got_df,
+                    f"Expected output for split {split} to be equal to the expected output.",
+                    check_column_order=False,
+                    check_row_order=False,
+                )
+
+                assert got_df["patient_id"].is_sorted(), f"Patient IDs should be sorted for split {split}."
+                for subj in splits[split]:
+                    got_df_subj = got_df.filter(pl.col("patient_id") == subj)
+                    assert got_df_subj[
+                        "time"
+                    ].is_sorted(), f"Times should be sorted for patient {subj} in split {split}."
+
+        except AssertionError as e:
+            print(f"Failed on split {split}")
+            print(f"stderr:\n{full_stderr}")
+            print(f"stdout:\n{full_stdout}")
+            raise e
+
+        # Stage 8: Finalize the metadata
+        stderr, stdout = run_command(
+            FINALIZE_METADATA_SCRIPT,
+            extraction_config_kwargs,
+            "finalize_metadata",
+        )
+        all_stderrs.append(stderr)
+        all_stdouts.append(stdout)
+
+        full_stderr = "\n".join(all_stderrs)
+        full_stdout = "\n".join(all_stdouts)
+
+        # Check code metadata
+        output_file = MEDS_cohort_dir / "metadata" / "codes.parquet"
+        assert output_file.is_file(), f"Expected {output_file} to exist: stderr:\n{stderr}\nstdout:\n{stdout}"
+
+        got_df = pl.read_parquet(output_file, glob=False, use_pyarrow=True)
+
+        want_df = pl.read_csv(source=StringIO(MEDS_OUTPUT_CODE_METADATA_FILE)).with_columns(
+            pl.col("code"),
+            pl.col("parent_codes").cast(pl.List(pl.Utf8)),
+        )
+
+        # We collapse the list type as it throws an error in the assert_df_equal otherwise
+        got_df = got_df.with_columns(pl.col("parent_codes").list.join("||"))
+        want_df = want_df.with_columns(pl.col("parent_codes").list.join("||"))
+
+        assert_df_equal(
+            want=want_df,
+            got=got_df,
+            msg=f"Finalized code metadata differs:\nstderr:\n{stderr}\nstdout:\n{stdout}",
+            check_column_order=False,
+            check_row_order=False,
+        )
+
+        # Check dataset metadata
+        output_file = MEDS_cohort_dir / "metadata" / "dataset.json"
+        assert output_file.is_file(), f"Expected {output_file} to exist: stderr:\n{stderr}\nstdout:\n{stdout}"
+
+        got_json = json.loads(output_file.read_text())
+        assert "etl_version" in got_json, "Expected 'etl_version' to be in the dataset metadata."
+        got_json.pop("etl_version")  # We don't test this as it changes with the commits.
+        assert got_json == MEDS_OUTPUT_DATASET_METADATA_JSON, f"Dataset metadata differs: {got_json}"
+
+        # Check the splits parquet
+        output_file = MEDS_cohort_dir / "metadata" / "patient_splits.parquet"
+        assert output_file.is_file(), f"Expected {output_file} to exist: stderr:\n{stderr}\nstdout:\n{stdout}"
+
+        got_df = pl.read_parquet(output_file, glob=False, use_pyarrow=True)
+        assert_df_equal(
+            PATIENT_SPLITS_DF,
+            got_df,
+            "Patient splits should be equal to the expected splits.",
+            check_column_order=False,
+            check_row_order=False,
+        )