Repository: mmcdermott/MEDS_transforms

Issue:
We should have a secondary multi-stage runner script that runs a full pipeline via a single command.
This runner script would take as input:
  1. A path to a traditional pipeline configuration file in yaml form (e.g., `preprocess.yaml`) stored on disk.
  2. A path to a secondary configuration file also on disk which maps stage names to run details, including:
    - script/command name
    - multirun type
    - hydra launcher kwargs (e.g., slurm parameters,number of workers, etc.)
  3. A command runner function (e.g., `subprocess`) which defaults to using python's typical method for launching command line commands. This argument is necessary for testing, so that we can test the runner without requiring it to run actual commands.

With these arguments, the runner would then simply, in sequence, step through and execute each stage with the appropriate arguments.

Discussion:

Patch:
diff --git a/pyproject.toml b/pyproject.toml
index c9f4906..756ac8e 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -24,7 +24,7 @@ dependencies = [
 
 [project.optional-dependencies]
 dev = ["pre-commit"]
-tests = ["pytest", "pytest-cov", "rootutils"]
+tests = ["pytest", "pytest-cov", "rootutils", "hydra-joblib-launcher"]
 local_parallelism = ["hydra-joblib-launcher"]
 slurm_parallelism = ["hydra-submitit-launcher"]
 docs = [
@@ -60,6 +60,9 @@ MEDS_transform-occlude_outliers = "MEDS_transforms.transforms.occlude_outliers:m
 MEDS_transform-tensorization = "MEDS_transforms.transforms.tensorization:main"
 MEDS_transform-tokenization = "MEDS_transforms.transforms.tokenization:main"
 
+# Runner
+MEDS_transform-runner = "MEDS_transforms.runner:main"
+
 [project.urls]
 Homepage = "https://github.com/mmcdermott/MEDS_transforms"
 Issues = "https://github.com/mmcdermott/MEDS_transforms/issues"
diff --git a/src/MEDS_transforms/__init__.py b/src/MEDS_transforms/__init__.py
index 8d0ffd6..2d62ae8 100644
--- a/src/MEDS_transforms/__init__.py
+++ b/src/MEDS_transforms/__init__.py
@@ -10,11 +10,14 @@ try:
 except PackageNotFoundError:  # pragma: no cover
     __version__ = "unknown"
 
-PREPROCESS_CONFIG_YAML = files(__package_name__).joinpath("configs/preprocess.yaml")
-EXTRACT_CONFIG_YAML = files(__package_name__).joinpath("configs/extract.yaml")
+PREPROCESS_CONFIG_YAML = files(__package_name__).joinpath("configs/_preprocess.yaml")
+EXTRACT_CONFIG_YAML = files(__package_name__).joinpath("configs/_extract.yaml")
+RUNNER_CONFIG_YAML = files(__package_name__).joinpath("configs/_runner.yaml")
 
 MANDATORY_COLUMNS = [subject_id_field, time_field, code_field, "numeric_value"]
 
+RESERVED_CONFIG_NAMES = {c.stem for c in (PREPROCESS_CONFIG_YAML, EXTRACT_CONFIG_YAML, RUNNER_CONFIG_YAML)}
+
 MANDATORY_TYPES = {
     subject_id_field: pl.Int64,
     time_field: pl.Datetime("us"),
diff --git a/src/MEDS_transforms/configs/extract.yaml b/src/MEDS_transforms/configs/_extract.yaml
similarity index 99%
rename from src/MEDS_transforms/configs/extract.yaml
rename to src/MEDS_transforms/configs/_extract.yaml
index 3abd498..dec483d 100644
--- a/src/MEDS_transforms/configs/extract.yaml
+++ b/src/MEDS_transforms/configs/_extract.yaml
@@ -1,5 +1,5 @@
 defaults:
-  - pipeline
+  - _pipeline
   - stage_configs:
       - shard_events
       - split_and_shard_subjects
diff --git a/src/MEDS_transforms/configs/pipeline.yaml b/src/MEDS_transforms/configs/_pipeline.yaml
similarity index 100%
rename from src/MEDS_transforms/configs/pipeline.yaml
rename to src/MEDS_transforms/configs/_pipeline.yaml
diff --git a/src/MEDS_transforms/configs/preprocess.yaml b/src/MEDS_transforms/configs/_preprocess.yaml
similarity index 98%
rename from src/MEDS_transforms/configs/preprocess.yaml
rename to src/MEDS_transforms/configs/_preprocess.yaml
index dab87a9..6ebafdc 100644
--- a/src/MEDS_transforms/configs/preprocess.yaml
+++ b/src/MEDS_transforms/configs/_preprocess.yaml
@@ -1,5 +1,5 @@
 defaults:
-  - pipeline
+  - _pipeline
   - stage_configs:
       - reshard_to_split
       - filter_subjects
diff --git a/src/MEDS_transforms/configs/_runner.yaml b/src/MEDS_transforms/configs/_runner.yaml
new file mode 100644
index 0000000..f826678
--- /dev/null
+++ b/src/MEDS_transforms/configs/_runner.yaml
@@ -0,0 +1,31 @@
+# Global IO
+pipeline_config_fp: ???
+stage_runner_fp: null
+
+_local_pipeline_config: ${oc.create:${load_yaml_file:${oc.select:pipeline_config_fp,null}}}
+_stage_runners: ${oc.create:${load_yaml_file:${stage_runner_fp}}}
+
+log_dir: "${_local_pipeline_config.cohort_dir}/.logs"
+
+_pipeline_description: ${oc.select:_local_pipeline_config.description,"No description provided."}
+
+do_profile: False
+
+hydra:
+  job:
+    name: "MEDS-transforms_runner_${now:%Y-%m-%d_%H-%M-%S}"
+  run:
+    dir: "${log_dir}"
+  help:
+    app_name: "MEDS-Transforms Pipeline Runner"
+
+    template: |-
+      == ${hydra.help.app_name} ==
+      ${hydra.help.app_name} is a command line tool for running entire MEDS-transform pipelines in a single
+      command.
+
+      ${get_script_docstring:runner}
+
+      **MEDS-transforms Pipeline description:**
+
+      ${_pipeline_description}
diff --git a/src/MEDS_transforms/mapreduce/mapper.py b/src/MEDS_transforms/mapreduce/mapper.py
index 6cc44e8..6247be6 100644
--- a/src/MEDS_transforms/mapreduce/mapper.py
+++ b/src/MEDS_transforms/mapreduce/mapper.py
@@ -621,11 +621,11 @@ def map_over(
     start = datetime.now()
 
     train_only = cfg.stage_cfg.get("train_only", False)
-    split_fp = Path(cfg.input_dir) / "metadata" / "subject_split.parquet"
 
     shards, includes_only_train = shard_iterator_fntr(cfg)
 
     if train_only:
+        split_fp = Path(cfg.input_dir) / "metadata" / "subject_splits.parquet"
         if includes_only_train:
             logger.info(
                 f"Processing train split only via shard prefix. Not filtering with {str(split_fp.resolve())}."
@@ -636,7 +636,7 @@ def map_over(
                 pl.scan_parquet(split_fp)
                 .filter(pl.col("split") == "train")
                 .select(subject_id_field)
-                .collect()
+                .collect()[subject_id_field]
                 .to_list()
             )
             read_fn = read_and_filter_fntr(train_subjects, read_fn)
diff --git a/src/MEDS_transforms/mapreduce/utils.py b/src/MEDS_transforms/mapreduce/utils.py
index a653d20..716ddc0 100644
--- a/src/MEDS_transforms/mapreduce/utils.py
+++ b/src/MEDS_transforms/mapreduce/utils.py
@@ -453,10 +453,11 @@ def shard_iterator(
         >>> includes_only_train
         False
 
-    If it can't find any files, it will return an empty list:
+    If it can't find any files, it will error:
         >>> fps, includes_only_train = shard_iterator(cfg)
-        >>> fps
-        []
+        Traceback (most recent call last):
+            ...
+        FileNotFoundError: No shards found in ... with suffix .parquet. Directory contents:...
     """
 
     input_dir = Path(cfg.stage_cfg.data_input_dir)
@@ -474,6 +475,12 @@ def shard_iterator(
         shard_name = shard_name[: -len(in_suffix)]
         shards.append(shard_name)
 
+    if not shards:
+        raise FileNotFoundError(
+            f"No shards found in {input_dir} with suffix {in_suffix}. Directory contents: "
+            f"{', '.join(str(p.relative_to(input_dir)) for p in input_dir.glob('**/*'))}"
+        )
+
     # We initialize this to False and overwrite it if we find dedicated train shards.
     includes_only_train = False
 
diff --git a/src/MEDS_transforms/runner.py b/src/MEDS_transforms/runner.py
new file mode 100755
index 0000000..8f82109
--- /dev/null
+++ b/src/MEDS_transforms/runner.py
@@ -0,0 +1,250 @@
+#!/usr/bin/env python
+"""This script is a helper utility to run entire pipelines from a single script.
+
+To do this effectively, this runner functionally takes a "meta configuration" file that contains:
+  1. The path to the pipeline configuration file.
+  2. Configuration details for how to run each stage of the pipeline, including mappings to the underlying
+     stage scripts and Hydra launcher configurations for each stage to control parallelism, resources, etc.
+"""
+
+import importlib
+import subprocess
+from pathlib import Path
+
+import hydra
+import yaml
+from loguru import logger
+from omegaconf import DictConfig, OmegaConf
+
+try:
+    from yaml import CLoader as Loader
+except ImportError:  # pragma: no cover
+    from yaml import Loader
+
+from MEDS_transforms import RESERVED_CONFIG_NAMES, RUNNER_CONFIG_YAML
+from MEDS_transforms.utils import hydra_loguru_init
+
+
+def get_script_from_name(stage_name: str) -> str | None:
+    """Returns the script name for the given stage name.
+
+    Args:
+        stage_name: The name of the stage.
+
+    Returns:
+        The script name for the given stage name.
+    """
+
+    try:
+        _ = importlib.import_module(f"MEDS_transforms.extract.{stage_name}")
+        return f"MEDS_extract-{stage_name}"
+    except ImportError:
+        pass
+
+    for pfx in ("MEDS_transforms.transforms", "MEDS_transforms.filters", "MEDS_transforms"):
+        try:
+            _ = importlib.import_module(f"{pfx}.{stage_name}")
+            return f"MEDS_transform-{stage_name}"
+        except ImportError:
+            pass
+
+    raise ValueError(f"Could not find a script for stage {stage_name}.")
+
+
+def get_parallelization_args(
+    parallelization_cfg: dict | DictConfig | None, default_parallelization_cfg: dict | DictConfig
+) -> list[str]:
+    """Gets the parallelization args."""
+
+    if parallelization_cfg is None:
+        return []
+
+    if len(parallelization_cfg) == 0 and len(default_parallelization_cfg) == 0:
+        return []
+
+    if "n_workers" in parallelization_cfg:
+        n_workers = parallelization_cfg["n_workers"]
+    elif "n_workers" in default_parallelization_cfg:
+        n_workers = default_parallelization_cfg["n_workers"]
+    else:
+        n_workers = 1
+
+    parallelization_args = [
+        "--multirun",
+        f'worker="range(0,{n_workers})"',
+    ]
+
+    if "launcher" in parallelization_cfg:
+        launcher = parallelization_cfg["launcher"]
+    elif "launcher" in default_parallelization_cfg:
+        launcher = default_parallelization_cfg["launcher"]
+    else:
+        launcher = None
+
+    if launcher is None:
+        return parallelization_args
+
+        if "launcher_params" in parallelization_cfg:
+            raise ValueError("If launcher_params is provided, launcher must also be provided.")
+
+    parallelization_args.append(f"hydra/launcher={launcher}")
+
+    if "launcher_params" in parallelization_cfg:
+        launcher_params = parallelization_cfg["launcher_params"]
+    elif "launcher_params" in default_parallelization_cfg:
+        launcher_params = default_parallelization_cfg["launcher_params"]
+    else:
+        launcher_params = {}
+
+    for k, v in launcher_params.items():
+        parallelization_args.append(f"hydra.launcher.{k}={v}")
+
+    return parallelization_args
+
+
+def run_stage(cfg: DictConfig, stage_name: str, default_parallelization_cfg: dict | DictConfig | None = None):
+    """Runs a single stage of the pipeline.
+
+    Args:
+        cfg: The configuration for the entire pipeline.
+        stage_name: The name of the stage to run.
+    """
+
+    if default_parallelization_cfg is None:
+        default_parallelization_cfg = {}
+
+    do_profile = cfg.get("do_profile", False)
+    pipeline_config_fp = Path(cfg.pipeline_config_fp)
+    stage_config = cfg._local_pipeline_config.stage_configs.get(stage_name, {})
+    stage_runner_config = cfg._stage_runners.get(stage_name, {})
+
+    script = None
+    if "script" in stage_runner_config:
+        script = stage_runner_config.script
+    elif "_script" in stage_config:
+        script = stage_config._script
+    else:
+        script = get_script_from_name(stage_name)
+
+    command_parts = [
+        script,
+        f"--config-dir={str(pipeline_config_fp.parent.resolve())}",
+        f"--config-name={pipeline_config_fp.stem}",
+        "'hydra.searchpath=[pkg://MEDS_transforms.configs]'",
+        f"stage={stage_name}",
+    ]
+
+    parallelization_args = get_parallelization_args(
+        stage_runner_config.get("parallelize", {}), default_parallelization_cfg
+    )
+
+    if parallelization_args:
+        multirun = parallelization_args.pop(0)
+        command_parts = command_parts[:3] + [multirun] + command_parts[3:] + parallelization_args
+
+    if do_profile:
+        command_parts.append("++hydra.callbacks.profiler._target_=hydra_profiler.profiler.ProfilerCallback")
+
+    full_cmd = " ".join(command_parts)
+    logger.info(f"Running command: {full_cmd}")
+    command_out = subprocess.run(full_cmd, shell=True, capture_output=True)
+
+    # https://stackoverflow.com/questions/21953835/run-subprocess-and-print-output-to-logging
+    # https://loguru.readthedocs.io/en/stable/api/logger.html#loguru._logger.Logger.parse
+
+    stderr = command_out.stderr.decode()
+    stdout = command_out.stdout.decode()
+    logger.info(f"Command output:\n{stdout}")
+    logger.info(f"Command error:\n{stderr}")
+
+    if command_out.returncode != 0:
+        raise ValueError(
+            f"Stage {stage_name} failed via {full_cmd} with return code {command_out.returncode}."
+        )
+
+
+@hydra.main(
+    version_base=None, config_path=str(RUNNER_CONFIG_YAML.parent), config_name=RUNNER_CONFIG_YAML.stem
+)
+def main(cfg: DictConfig):
+    """Runs the entire pipeline, end-to-end, based on the configuration provided.
+
+    This script will launch many subsidiary commands via `subprocess`, one for each stage of the specified
+    pipeline.
+    """
+
+    hydra_loguru_init()
+
+    pipeline_config_fp = Path(cfg.pipeline_config_fp)
+    if not pipeline_config_fp.exists():
+        raise FileNotFoundError(f"Pipeline configuration file {pipeline_config_fp} does not exist.")
+    if not pipeline_config_fp.suffix == ".yaml":
+        raise ValueError(f"Pipeline configuration file {pipeline_config_fp} must have a .yaml extension.")
+    if pipeline_config_fp.stem in RESERVED_CONFIG_NAMES:
+        raise ValueError(
+            f"Pipeline configuration file {pipeline_config_fp} must not have a name in "
+            f"{RESERVED_CONFIG_NAMES}."
+        )
+
+    pipeline_config = load_yaml_file(cfg.pipeline_config_fp)
+    stages = pipeline_config.get("stages", [])
+    if not stages:
+        raise ValueError("Pipeline configuration must specify at least one stage.")
+
+    log_dir = Path(cfg.log_dir)
+
+    if cfg.get("do_profile", False):
+        try:
+            import hydra_profiler  # noqa: F401
+        except ImportError as e:
+            raise ValueError(
+                "You can't run in profiling mode without installing hydra-profiler. Try installing "
+                "MEDS-transforms with the 'profiler' optional dependency: "
+                "`pip install MEDS-transforms[profiler]`."
+            ) from e
+
+    global_done_file = log_dir / "_all_stages.done"
+    if global_done_file.exists():
+        logger.info("All stages are already complete. Exiting.")
+        return
+
+    if "parallelize" in cfg._stage_runners:
+        default_parallelization_cfg = cfg._stage_runners.parallelize
+    elif "parallelize" in cfg:
+        default_parallelization_cfg = cfg.parallelize
+    else:
+        default_parallelization_cfg = None
+
+    for stage in stages:
+        done_file = log_dir / f"{stage}.done"
+
+        if done_file.exists():
+            logger.info(f"Skipping stage {stage} as it is already complete.")
+        else:
+            logger.info(f"Running stage: {stage}")
+            run_stage(cfg, stage, default_parallelization_cfg=default_parallelization_cfg)
+            done_file.touch()
+
+    global_done_file.touch()
+
+
+def load_yaml_file(path: str | None) -> dict | DictConfig:
+    if not path:
+        return {}
+
+    path = Path(path)
+    if not path.exists():
+        raise FileNotFoundError(f"File {path} does not exist.")
+
+    try:
+        return OmegaConf.load(path)
+    except Exception as e:
+        logger.warning(f"Failed to load {path} as an OmegaConf: {e}. Trying as a plain YAML file.")
+        yaml_text = path.read_text()
+        return yaml.load(yaml_text, Loader=Loader)
+
+
+OmegaConf.register_new_resolver("load_yaml_file", load_yaml_file, replace=True)
+
+if __name__ == "__main__":  # pragma: no cover
+    main()
diff --git a/src/MEDS_transforms/utils.py b/src/MEDS_transforms/utils.py
index b62f7d1..871b90a 100644
--- a/src/MEDS_transforms/utils.py
+++ b/src/MEDS_transforms/utils.py
@@ -1,5 +1,6 @@
 """Core utilities for MEDS pipelines built with these tools."""
 
+import importlib
 import inspect
 import os
 import sys
@@ -108,10 +109,13 @@ def get_package_version() -> str:
     return package_version
 
 
-def get_script_docstring() -> str:
+def get_script_docstring(filename: str | None = None) -> str:
     """Returns the docstring of the main function of the script from which this function was called."""
 
-    main_module = sys.modules["__main__"]
+    if filename is not None:
+        main_module = importlib.import_module(f"MEDS_transforms.{filename}")
+    else:
+        main_module = sys.modules["__main__"]
     func = getattr(main_module, "main", None)
     if func and callable(func):
         return inspect.getdoc(func) or ""


Test patch:
diff --git a/tests/MEDS_Extract/test_convert_to_sharded_events.py b/tests/MEDS_Extract/test_convert_to_sharded_events.py
index 074e897..6d8311d 100644
--- a/tests/MEDS_Extract/test_convert_to_sharded_events.py
+++ b/tests/MEDS_Extract/test_convert_to_sharded_events.py
@@ -104,107 +104,107 @@ SHARDS_JSON = {
 
 WANT_OUTPUTS = parse_shards_yaml(
     """
-    data/train/0/subjects/[0-6).parquet: |-2
-      subject_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,
-
-    data/train/1/subjects/[0-6).parquet: |-2
-      subject_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,
-
-    data/tuning/0/subjects/[0-6).parquet: |-2
-      subject_id,time,code,numeric_value
-      754281,,EYE_COLOR//BROWN,
-      754281,,HEIGHT,166.22261567137025
-      754281,"12/19/1988, 00:00:00",DOB,
-
-    data/held_out/0/subjects/[0-6).parquet: |-2
-      subject_id,time,code,numeric_value
-      1500733,,EYE_COLOR//BROWN,
-      1500733,,HEIGHT,158.60131573580904
-      1500733,"07/20/1986, 00:00:00",DOB,
-
-    data/train/0/admit_vitals/[0-10).parquet: |-2
-      subject_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:25:32",HR,114.1
-      1195293,"06/20/2010, 19:25:32",TEMP,100.0
-      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:50:04",DISCHARGE,
-
-    data/train/0/admit_vitals/[10-16).parquet: |-2
-      subject_id,time,code,numeric_value
-      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:45:19",HR,119.8
-      1195293,"06/20/2010, 19:45:19",TEMP,99.9
-      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,
-
-    data/train/1/admit_vitals/[0-10).parquet: |-2
-      subject_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,
-
-    data/train/1/admit_vitals/[10-16).parquet: |-2
-      subject_id,time,code,numeric_value
-
-    data/tuning/0/admit_vitals/[0-10).parquet: |-2
-      subject_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,
-
-    data/tuning/0/admit_vitals/[10-16).parquet: |-2
-      subject_id,time,code,numeric_value
-
-    data/held_out/0/admit_vitals/[0-10).parquet: |-2
-      subject_id,time,code,numeric_value
-      1500733,"06/03/2010, 14:54:38",ADMISSION//ORTHOPEDIC,
-      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,
-
-    data/held_out/0/admit_vitals/[10-16).parquet: |-2
-      subject_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:44:26",DISCHARGE,
+data/train/0/subjects/[0-6).parquet: |-2
+  subject_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,
+
+data/train/1/subjects/[0-6).parquet: |-2
+  subject_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,
+
+data/tuning/0/subjects/[0-6).parquet: |-2
+  subject_id,time,code,numeric_value
+  754281,,EYE_COLOR//BROWN,
+  754281,,HEIGHT,166.22261567137025
+  754281,"12/19/1988, 00:00:00",DOB,
+
+data/held_out/0/subjects/[0-6).parquet: |-2
+  subject_id,time,code,numeric_value
+  1500733,,EYE_COLOR//BROWN,
+  1500733,,HEIGHT,158.60131573580904
+  1500733,"07/20/1986, 00:00:00",DOB,
+
+data/train/0/admit_vitals/[0-10).parquet: |-2
+  subject_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:25:32",HR,114.1
+  1195293,"06/20/2010, 19:25:32",TEMP,100.0
+  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:50:04",DISCHARGE,
+
+data/train/0/admit_vitals/[10-16).parquet: |-2
+  subject_id,time,code,numeric_value
+  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:45:19",HR,119.8
+  1195293,"06/20/2010, 19:45:19",TEMP,99.9
+  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,
+
+data/train/1/admit_vitals/[0-10).parquet: |-2
+  subject_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,
+
+data/train/1/admit_vitals/[10-16).parquet: |-2
+  subject_id,time,code,numeric_value
+
+data/tuning/0/admit_vitals/[0-10).parquet: |-2
+  subject_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,
+
+data/tuning/0/admit_vitals/[10-16).parquet: |-2
+  subject_id,time,code,numeric_value
+
+data/held_out/0/admit_vitals/[0-10).parquet: |-2
+  subject_id,time,code,numeric_value
+  1500733,"06/03/2010, 14:54:38",ADMISSION//ORTHOPEDIC,
+  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,
+
+data/held_out/0/admit_vitals/[10-16).parquet: |-2
+  subject_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:44:26",DISCHARGE,
     """
 )
 
diff --git a/tests/MEDS_Extract/test_extract_code_metadata.py b/tests/MEDS_Extract/test_extract_code_metadata.py
index e93bc34..7700426 100644
--- a/tests/MEDS_Extract/test_extract_code_metadata.py
+++ b/tests/MEDS_Extract/test_extract_code_metadata.py
@@ -12,80 +12,80 @@ from tests.utils import parse_shards_yaml, single_stage_tester
 
 INPUT_SHARDS = parse_shards_yaml(
     """
-    data/train/0: |-2
-      subject_id,time,code,numeric_value
-      239684,,EYE_COLOR//BROWN,
-      239684,,HEIGHT,175.271115221764
-      239684,"12/28/1980, 00:00:00",DOB,
-      239684,"05/11/2010, 17:41:51",ADMISSION//CARDIAC,
-      239684,"05/11/2010, 17:41:51",HR//bpm,102.6
-      239684,"05/11/2010, 17:41:51",TEMP,96.0
-      239684,"05/11/2010, 17:48:48",HR//bpm,105.1
-      239684,"05/11/2010, 17:48:48",TEMP,96.2
-      239684,"05/11/2010, 18:25:35",HR//bpm,113.4
-      239684,"05/11/2010, 18:25:35",TEMP,95.8
-      239684,"05/11/2010, 18:57:18",HR//bpm,112.6
-      239684,"05/11/2010, 18:57:18",TEMP,95.5
-      239684,"05/11/2010, 19:27:19",DISCHARGE,
-      1195293,,EYE_COLOR//BLUE,
-      1195293,,HEIGHT,164.6868838269085
-      1195293,"06/20/1978, 00:00:00",DOB,
-      1195293,"06/20/2010, 19:23:52",ADMISSION//CARDIAC,
-      1195293,"06/20/2010, 19:23:52",HR//bpm,109.0
-      1195293,"06/20/2010, 19:23:52",TEMP,100.0
-      1195293,"06/20/2010, 19:25:32",HR//bpm,114.1
-      1195293,"06/20/2010, 19:25:32",TEMP,100.0
-      1195293,"06/20/2010, 19:45:19",HR//bpm,119.8
-      1195293,"06/20/2010, 19:45:19",TEMP,99.9
-      1195293,"06/20/2010, 20:12:31",HR//bpm,112.5
-      1195293,"06/20/2010, 20:12:31",TEMP,99.8
-      1195293,"06/20/2010, 20:24:44",HR//bpm,107.7
-      1195293,"06/20/2010, 20:24:44",TEMP,100.0
-      1195293,"06/20/2010, 20:41:33",HR//bpm,107.5
-      1195293,"06/20/2010, 20:41:33",TEMP,100.4
-      1195293,"06/20/2010, 20:50:04",DISCHARGE,
-
-
-    data/train/1: |-2
-      subject_id,time,code,numeric_value
-      68729,,EYE_COLOR//HAZEL,
-      68729,,HEIGHT,160.3953106166676
-      68729,"03/09/1978, 00:00:00",DOB,
-      68729,"05/26/2010, 02:30:56",ADMISSION//PULMONARY,
-      68729,"05/26/2010, 02:30:56",HR//bpm,86.0
-      68729,"05/26/2010, 02:30:56",TEMP,97.8
-      68729,"05/26/2010, 04:51:52",DISCHARGE,
-      814703,,EYE_COLOR//HAZEL,
-      814703,,HEIGHT,156.48559093209357
-      814703,"03/28/1976, 00:00:00",DOB,
-      814703,"02/05/2010, 05:55:39",ADMISSION//ORTHOPEDIC,
-      814703,"02/05/2010, 05:55:39",HR//bpm,170.2
-      814703,"02/05/2010, 05:55:39",TEMP,100.1
-      814703,"02/05/2010, 07:02:30",DISCHARGE,
-
-    data/tuning/0: |-2
-      subject_id,time,code,numeric_value
-      754281,,EYE_COLOR//BROWN,
-      754281,,HEIGHT,166.22261567137025
-      754281,"12/19/1988, 00:00:00",DOB,
-      754281,"01/03/2010, 06:27:59",ADMISSION//PULMONARY,
-      754281,"01/03/2010, 06:27:59",HR//bpm,142.0
-      754281,"01/03/2010, 06:27:59",TEMP,99.8
-      754281,"01/03/2010, 08:22:13",DISCHARGE,
-
-    data/held_out/0: |-2
-      subject_id,time,code,numeric_value
-      1500733,,EYE_COLOR//BROWN,
-      1500733,,HEIGHT,158.60131573580904
-      1500733,"07/20/1986, 00:00:00",DOB,
-      1500733,"06/03/2010, 14:54:38",ADMISSION//ORTHOPEDIC,
-      1500733,"06/03/2010, 14:54:38",HR//bpm,91.4
-      1500733,"06/03/2010, 14:54:38",TEMP,100.0
-      1500733,"06/03/2010, 15:39:49",HR//bpm,84.4
-      1500733,"06/03/2010, 15:39:49",TEMP,100.3
-      1500733,"06/03/2010, 16:20:49",HR//bpm,90.1
-      1500733,"06/03/2010, 16:20:49",TEMP,100.1
-      1500733,"06/03/2010, 16:44:26",DISCHARGE,
+data/train/0: |-2
+  subject_id,time,code,numeric_value
+  239684,,EYE_COLOR//BROWN,
+  239684,,HEIGHT,175.271115221764
+  239684,"12/28/1980, 00:00:00",DOB,
+  239684,"05/11/2010, 17:41:51",ADMISSION//CARDIAC,
+  239684,"05/11/2010, 17:41:51",HR//bpm,102.6
+  239684,"05/11/2010, 17:41:51",TEMP,96.0
+  239684,"05/11/2010, 17:48:48",HR//bpm,105.1
+  239684,"05/11/2010, 17:48:48",TEMP,96.2
+  239684,"05/11/2010, 18:25:35",HR//bpm,113.4
+  239684,"05/11/2010, 18:25:35",TEMP,95.8
+  239684,"05/11/2010, 18:57:18",HR//bpm,112.6
+  239684,"05/11/2010, 18:57:18",TEMP,95.5
+  239684,"05/11/2010, 19:27:19",DISCHARGE,
+  1195293,,EYE_COLOR//BLUE,
+  1195293,,HEIGHT,164.6868838269085
+  1195293,"06/20/1978, 00:00:00",DOB,
+  1195293,"06/20/2010, 19:23:52",ADMISSION//CARDIAC,
+  1195293,"06/20/2010, 19:23:52",HR//bpm,109.0
+  1195293,"06/20/2010, 19:23:52",TEMP,100.0
+  1195293,"06/20/2010, 19:25:32",HR//bpm,114.1
+  1195293,"06/20/2010, 19:25:32",TEMP,100.0
+  1195293,"06/20/2010, 19:45:19",HR//bpm,119.8
+  1195293,"06/20/2010, 19:45:19",TEMP,99.9
+  1195293,"06/20/2010, 20:12:31",HR//bpm,112.5
+  1195293,"06/20/2010, 20:12:31",TEMP,99.8
+  1195293,"06/20/2010, 20:24:44",HR//bpm,107.7
+  1195293,"06/20/2010, 20:24:44",TEMP,100.0
+  1195293,"06/20/2010, 20:41:33",HR//bpm,107.5
+  1195293,"06/20/2010, 20:41:33",TEMP,100.4
+  1195293,"06/20/2010, 20:50:04",DISCHARGE,
+
+
+data/train/1: |-2
+  subject_id,time,code,numeric_value
+  68729,,EYE_COLOR//HAZEL,
+  68729,,HEIGHT,160.3953106166676
+  68729,"03/09/1978, 00:00:00",DOB,
+  68729,"05/26/2010, 02:30:56",ADMISSION//PULMONARY,
+  68729,"05/26/2010, 02:30:56",HR//bpm,86.0
+  68729,"05/26/2010, 02:30:56",TEMP,97.8
+  68729,"05/26/2010, 04:51:52",DISCHARGE,
+  814703,,EYE_COLOR//HAZEL,
+  814703,,HEIGHT,156.48559093209357
+  814703,"03/28/1976, 00:00:00",DOB,
+  814703,"02/05/2010, 05:55:39",ADMISSION//ORTHOPEDIC,
+  814703,"02/05/2010, 05:55:39",HR//bpm,170.2
+  814703,"02/05/2010, 05:55:39",TEMP,100.1
+  814703,"02/05/2010, 07:02:30",DISCHARGE,
+
+data/tuning/0: |-2
+  subject_id,time,code,numeric_value
+  754281,,EYE_COLOR//BROWN,
+  754281,,HEIGHT,166.22261567137025
+  754281,"12/19/1988, 00:00:00",DOB,
+  754281,"01/03/2010, 06:27:59",ADMISSION//PULMONARY,
+  754281,"01/03/2010, 06:27:59",HR//bpm,142.0
+  754281,"01/03/2010, 06:27:59",TEMP,99.8
+  754281,"01/03/2010, 08:22:13",DISCHARGE,
+
+data/held_out/0: |-2
+  subject_id,time,code,numeric_value
+  1500733,,EYE_COLOR//BROWN,
+  1500733,,HEIGHT,158.60131573580904
+  1500733,"07/20/1986, 00:00:00",DOB,
+  1500733,"06/03/2010, 14:54:38",ADMISSION//ORTHOPEDIC,
+  1500733,"06/03/2010, 14:54:38",HR//bpm,91.4
+  1500733,"06/03/2010, 14:54:38",TEMP,100.0
+  1500733,"06/03/2010, 15:39:49",HR//bpm,84.4
+  1500733,"06/03/2010, 15:39:49",TEMP,100.3
+  1500733,"06/03/2010, 16:20:49",HR//bpm,90.1
+  1500733,"06/03/2010, 16:20:49",TEMP,100.1
+  1500733,"06/03/2010, 16:44:26",DISCHARGE,
     """
 )
 
diff --git a/tests/MEDS_Extract/test_finalize_MEDS_data.py b/tests/MEDS_Extract/test_finalize_MEDS_data.py
index d9a3e0a..3348d12 100644
--- a/tests/MEDS_Extract/test_finalize_MEDS_data.py
+++ b/tests/MEDS_Extract/test_finalize_MEDS_data.py
@@ -11,81 +11,81 @@ from tests.utils import parse_shards_yaml, single_stage_tester
 
 INPUT_SHARDS = parse_shards_yaml(
     """
-    data/train/0: |-2
-      subject_id,time,code,numeric_value
-      239684,,EYE_COLOR//BROWN,
-      239684,,HEIGHT,175.271115221764
-      239684,"12/28/1980, 00:00:00",DOB,
-      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,,EYE_COLOR//BLUE,
-      1195293,,HEIGHT,164.6868838269085
-      1195293,"06/20/1978, 00:00:00",DOB,
-      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,
+data/train/0: |-2
+  subject_id,time,code,numeric_value
+  239684,,EYE_COLOR//BROWN,
+  239684,,HEIGHT,175.271115221764
+  239684,"12/28/1980, 00:00:00",DOB,
+  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,,EYE_COLOR//BLUE,
+  1195293,,HEIGHT,164.6868838269085
+  1195293,"06/20/1978, 00:00:00",DOB,
+  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,
 
 
-    data/train/1: |-2
-      subject_id,time,code,numeric_value
-      68729,,EYE_COLOR//HAZEL,
-      68729,,HEIGHT,160.3953106166676
-      68729,"03/09/1978, 00:00:00",DOB,
-      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,,EYE_COLOR//HAZEL,
-      814703,,HEIGHT,156.48559093209357
-      814703,"03/28/1976, 00:00:00",DOB,
-      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,
+data/train/1: |-2
+  subject_id,time,code,numeric_value
+  68729,,EYE_COLOR//HAZEL,
+  68729,,HEIGHT,160.3953106166676
+  68729,"03/09/1978, 00:00:00",DOB,
+  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,,EYE_COLOR//HAZEL,
+  814703,,HEIGHT,156.48559093209357
+  814703,"03/28/1976, 00:00:00",DOB,
+  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,
 
-    data/tuning/0: |-2
-      subject_id,time,code,numeric_value
-      754281,,EYE_COLOR//BROWN,
-      754281,,HEIGHT,166.22261567137025
-      754281,"12/19/1988, 00:00:00",DOB,
-      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,
+data/tuning/0: |-2
+  subject_id,time,code,numeric_value
+  754281,,EYE_COLOR//BROWN,
+  754281,,HEIGHT,166.22261567137025
+  754281,"12/19/1988, 00:00:00",DOB,
+  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,
 
-    data/held_out/0: |-2
-      subject_id,time,code,numeric_value
-      1500733,,EYE_COLOR//BROWN,
-      1500733,,HEIGHT,158.60131573580904
-      1500733,"07/20/1986, 00:00:00",DOB,
-      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,
-    """
+data/held_out/0: |-2
+  subject_id,time,code,numeric_value
+  1500733,,EYE_COLOR//BROWN,
+  1500733,,HEIGHT,158.60131573580904
+  1500733,"07/20/1986, 00:00:00",DOB,
+  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,
+"""
 )
 
 WANT_OUTPUTS = {
diff --git a/tests/MEDS_Extract/test_merge_to_MEDS_cohort.py b/tests/MEDS_Extract/test_merge_to_MEDS_cohort.py
index b9a21ff..7468804 100644
--- a/tests/MEDS_Extract/test_merge_to_MEDS_cohort.py
+++ b/tests/MEDS_Extract/test_merge_to_MEDS_cohort.py
@@ -66,186 +66,186 @@ SHARDS_JSON = {
 
 INPUT_SHARDS = parse_shards_yaml(
     """
-    data/train/0/subjects/[0-6): |-2
-      subject_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,
+data/train/0/subjects/[0-6): |-2
+  subject_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,
 
-    data/train/1/subjects/[0-6): |-2
-      subject_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,
+data/train/1/subjects/[0-6): |-2
+  subject_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,
 
-    data/tuning/0/subjects/[0-6): |-2
-      subject_id,time,code,numeric_value
-      754281,,EYE_COLOR//BROWN,
-      754281,,HEIGHT,166.22261567137025
-      754281,"12/19/1988, 00:00:00",DOB,
+data/tuning/0/subjects/[0-6): |-2
+  subject_id,time,code,numeric_value
+  754281,,EYE_COLOR//BROWN,
+  754281,,HEIGHT,166.22261567137025
+  754281,"12/19/1988, 00:00:00",DOB,
 
-    data/held_out/0/subjects/[0-6): |-2
-      subject_id,time,code,numeric_value
-      1500733,,EYE_COLOR//BROWN,
-      1500733,,HEIGHT,158.60131573580904
-      1500733,"07/20/1986, 00:00:00",DOB,
+data/held_out/0/subjects/[0-6): |-2
+  subject_id,time,code,numeric_value
+  1500733,,EYE_COLOR//BROWN,
+  1500733,,HEIGHT,158.60131573580904
+  1500733,"07/20/1986, 00:00:00",DOB,
 
-    data/train/0/admit_vitals/[0-10): |-2
-      subject_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:25:32",HR,114.1
-      1195293,"06/20/2010, 19:25:32",TEMP,100.0
-      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:50:04",DISCHARGE,
+data/train/0/admit_vitals/[0-10): |-2
+  subject_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:25:32",HR,114.1
+  1195293,"06/20/2010, 19:25:32",TEMP,100.0
+  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:50:04",DISCHARGE,
 
-    data/train/0/admit_vitals/[10-16): |-2
-      subject_id,time,code,numeric_value
-      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:45:19",HR,119.8
-      1195293,"06/20/2010, 19:45:19",TEMP,99.9
-      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,
+data/train/0/admit_vitals/[10-16): |-2
+  subject_id,time,code,numeric_value
+  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:45:19",HR,119.8
+  1195293,"06/20/2010, 19:45:19",TEMP,99.9
+  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,
 
-    data/train/1/admit_vitals/[0-10): |-2
-      subject_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,
+data/train/1/admit_vitals/[0-10): |-2
+  subject_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,
 
-    data/train/1/admit_vitals/[10-16): |-2
-      subject_id,time,code,numeric_value
+data/train/1/admit_vitals/[10-16): |-2
+  subject_id,time,code,numeric_value
 
-    data/tuning/0/admit_vitals/[0-10): |-2
-      subject_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,
+data/tuning/0/admit_vitals/[0-10): |-2
+  subject_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,
 
-    data/tuning/0/admit_vitals/[10-16): |-2
-      subject_id,time,code,numeric_value
+data/tuning/0/admit_vitals/[10-16): |-2
+  subject_id,time,code,numeric_value
 
-    data/held_out/0/admit_vitals/[0-10): |-2
-      subject_id,time,code,numeric_value
-      1500733,"06/03/2010, 14:54:38",ADMISSION//ORTHOPEDIC,
-      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,
+data/held_out/0/admit_vitals/[0-10): |-2
+  subject_id,time,code,numeric_value
+  1500733,"06/03/2010, 14:54:38",ADMISSION//ORTHOPEDIC,
+  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,
 
-    data/held_out/0/admit_vitals/[10-16): |-2
-      subject_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:44:26",DISCHARGE,
+data/held_out/0/admit_vitals/[10-16): |-2
+  subject_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:44:26",DISCHARGE,
     """
 )
 
 WANT_OUTPUTS = parse_shards_yaml(
     """
-    data/train/0: |-2
-      subject_id,time,code,numeric_value
-      239684,,EYE_COLOR//BROWN,
-      239684,,HEIGHT,175.271115221764
-      239684,"12/28/1980, 00:00:00",DOB,
-      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,,EYE_COLOR//BLUE,
-      1195293,,HEIGHT,164.6868838269085
-      1195293,"06/20/1978, 00:00:00",DOB,
-      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,
+data/train/0: |-2
+  subject_id,time,code,numeric_value
+  239684,,EYE_COLOR//BROWN,
+  239684,,HEIGHT,175.271115221764
+  239684,"12/28/1980, 00:00:00",DOB,
+  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,,EYE_COLOR//BLUE,
+  1195293,,HEIGHT,164.6868838269085
+  1195293,"06/20/1978, 00:00:00",DOB,
+  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,
 
 
-    data/train/1: |-2
-      subject_id,time,code,numeric_value
-      68729,,EYE_COLOR//HAZEL,
-      68729,,HEIGHT,160.3953106166676
-      68729,"03/09/1978, 00:00:00",DOB,
-      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,,EYE_COLOR//HAZEL,
-      814703,,HEIGHT,156.48559093209357
-      814703,"03/28/1976, 00:00:00",DOB,
-      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,
+data/train/1: |-2
+  subject_id,time,code,numeric_value
+  68729,,EYE_COLOR//HAZEL,
+  68729,,HEIGHT,160.3953106166676
+  68729,"03/09/1978, 00:00:00",DOB,
+  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,,EYE_COLOR//HAZEL,
+  814703,,HEIGHT,156.48559093209357
+  814703,"03/28/1976, 00:00:00",DOB,
+  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,
 
-    data/tuning/0: |-2
-      subject_id,time,code,numeric_value
-      754281,,EYE_COLOR//BROWN,
-      754281,,HEIGHT,166.22261567137025
-      754281,"12/19/1988, 00:00:00",DOB,
-      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,
+data/tuning/0: |-2
+  subject_id,time,code,numeric_value
+  754281,,EYE_COLOR//BROWN,
+  754281,,HEIGHT,166.22261567137025
+  754281,"12/19/1988, 00:00:00",DOB,
+  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,
 
-    data/held_out/0: |-2
-      subject_id,time,code,numeric_value
-      1500733,,EYE_COLOR//BROWN,
-      1500733,,HEIGHT,158.60131573580904
-      1500733,"07/20/1986, 00:00:00",DOB,
-      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,
+data/held_out/0: |-2
+  subject_id,time,code,numeric_value
+  1500733,,EYE_COLOR//BROWN,
+  1500733,,HEIGHT,158.60131573580904
+  1500733,"07/20/1986, 00:00:00",DOB,
+  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,
     """
 )
 
diff --git a/tests/MEDS_Transforms/test_extract_values.py b/tests/MEDS_Transforms/test_extract_values.py
index 4114b3b..83a2aa3 100644
--- a/tests/MEDS_Transforms/test_extract_values.py
+++ b/tests/MEDS_Transforms/test_extract_values.py
@@ -9,71 +9,71 @@ from tests.MEDS_Transforms.transform_tester_base import parse_shards_yaml, singl
 
 INPUT_SHARDS = parse_shards_yaml(
     """
-    train/0: |-2
-      subject_id,time,code,numeric_value,text_value
-      239684,,EYE_COLOR//BROWN,,
-      239684,"12/28/1980, 00:00:00",DOB,,
-      239684,"05/11/2010, 17:41:51",BP,,"120/80"
-      1195293,,EYE_COLOR//BLUE,,
-      1195293,"06/20/1978, 00:00:00",DOB,,
-      1195293,"06/20/2010, 19:23:52",BP,,"144/96"
-      1195293,"06/20/2010, 19:23:52",HR,80,
-      1195293,"06/20/2010, 19:23:52",TEMP,,"100F"
-    train/1: |-2
-      subject_id,time,code,numeric_value,text_value
-      68729,,EYE_COLOR//HAZEL,,
-      68729,"03/09/1978, 00:00:00",DOB,,
-      814703,"02/05/2010, 05:55:39",HR,170.2,
-    tuning/0: |-2
-      subject_id,time,code,numeric_value,text_value
-      754281,,EYE_COLOR//BROWN,,
-      754281,"12/19/1988, 00:00:00",DOB,,
-      754281,"01/03/2010, 06:27:59",HR,142.0,
-      754281,"06/20/2010, 20:23:50",BP,,"134/76"
-      754281,"06/20/2010, 21:00:02",TEMP,,"36.2C"
-    held_out/0: |-2
-      subject_id,time,code,numeric_value,text_value
-      1500733,,EYE_COLOR//BROWN,,
-      1500733,"07/20/1986, 00:00:00",DOB,,
-      1500733,"06/03/2010, 14:54:38",HR,91.4
-      1500733,"06/03/2010, 14:54:38",BP,,"123/82"
+train/0: |-2
+  subject_id,time,code,numeric_value,text_value
+  239684,,EYE_COLOR//BROWN,,
+  239684,"12/28/1980, 00:00:00",DOB,,
+  239684,"05/11/2010, 17:41:51",BP,,"120/80"
+  1195293,,EYE_COLOR//BLUE,,
+  1195293,"06/20/1978, 00:00:00",DOB,,
+  1195293,"06/20/2010, 19:23:52",BP,,"144/96"
+  1195293,"06/20/2010, 19:23:52",HR,80,
+  1195293,"06/20/2010, 19:23:52",TEMP,,"100F"
+train/1: |-2
+  subject_id,time,code,numeric_value,text_value
+  68729,,EYE_COLOR//HAZEL,,
+  68729,"03/09/1978, 00:00:00",DOB,,
+  814703,"02/05/2010, 05:55:39",HR,170.2,
+tuning/0: |-2
+  subject_id,time,code,numeric_value,text_value
+  754281,,EYE_COLOR//BROWN,,
+  754281,"12/19/1988, 00:00:00",DOB,,
+  754281,"01/03/2010, 06:27:59",HR,142.0,
+  754281,"06/20/2010, 20:23:50",BP,,"134/76"
+  754281,"06/20/2010, 21:00:02",TEMP,,"36.2C"
+held_out/0: |-2
+  subject_id,time,code,numeric_value,text_value
+  1500733,,EYE_COLOR//BROWN,,
+  1500733,"07/20/1986, 00:00:00",DOB,,
+  1500733,"06/03/2010, 14:54:38",HR,91.4
+  1500733,"06/03/2010, 14:54:38",BP,,"123/82"
     """
 )
 
 WANT_SHARDS = parse_shards_yaml(
     """
-    train/0: |-2
-      subject_id,time,code,numeric_value,text_value
-      239684,,EYE_COLOR//BROWN,,
-      239684,"12/28/1980, 00:00:00",DOB,,
-      239684,"05/11/2010, 17:41:51",BP//SYSTOLIC,120,
-      239684,"05/11/2010, 17:41:51",BP//DIASTOLIC,80,
-      1195293,,EYE_COLOR//BLUE,,
-      1195293,"06/20/1978, 00:00:00",DOB,,
-      1195293,"06/20/2010, 19:23:52",BP//SYSTOLIC,144,
-      1195293,"06/20/2010, 19:23:52",BP//DIASTOLIC,96,
-      1195293,"06/20/2010, 19:23:52",TEMP//F,100,
-      1195293,"06/20/2010, 19:23:52",HR,80,
-    train/1: |-2
-      subject_id,time,code,numeric_value,text_value
-      68729,,EYE_COLOR//HAZEL,,
-      68729,"03/09/1978, 00:00:00",DOB,,
-      814703,"02/05/2010, 05:55:39",HR,170.2,
-    tuning/0: |-2
-      subject_id,time,code,numeric_value,text_value
-      754281,,EYE_COLOR//BROWN,,
-      754281,"12/19/1988, 00:00:00",DOB,,
-      754281,"01/03/2010, 06:27:59",HR,142.0,
-      754281,"06/20/2010, 20:23:50",BP//SYSTOLIC,134,
-      754281,"06/20/2010, 20:23:50",BP//DIASTOLIC,76,
-      754281,"06/20/2010, 21:00:02",TEMP//C,36.2,
-    held_out/0: |-2
-      subject_id,time,code,numeric_value,text_value
-      1500733,,EYE_COLOR//BROWN,,
-      1500733,"07/20/1986, 00:00:00",DOB,,
-      1500733,"06/03/2010, 14:54:38",BP//SYSTOLIC,123,
-      1500733,"06/03/2010, 14:54:38",BP//DIASTOLIC,82,
-      1500733,"06/03/2010, 14:54:38",HR,91.4,
+train/0: |-2
+  subject_id,time,code,numeric_value,text_value
+  239684,,EYE_COLOR//BROWN,,
+  239684,"12/28/1980, 00:00:00",DOB,,
+  239684,"05/11/2010, 17:41:51",BP//SYSTOLIC,120,
+  239684,"05/11/2010, 17:41:51",BP//DIASTOLIC,80,
+  1195293,,EYE_COLOR//BLUE,,
+  1195293,"06/20/1978, 00:00:00",DOB,,
+  1195293,"06/20/2010, 19:23:52",BP//SYSTOLIC,144,
+  1195293,"06/20/2010, 19:23:52",BP//DIASTOLIC,96,
+  1195293,"06/20/2010, 19:23:52",TEMP//F,100,
+  1195293,"06/20/2010, 19:23:52",HR,80,
+train/1: |-2
+  subject_id,time,code,numeric_value,text_value
+  68729,,EYE_COLOR//HAZEL,,
+  68729,"03/09/1978, 00:00:00",DOB,,
+  814703,"02/05/2010, 05:55:39",HR,170.2,
+tuning/0: |-2
+  subject_id,time,code,numeric_value,text_value
+  754281,,EYE_COLOR//BROWN,,
+  754281,"12/19/1988, 00:00:00",DOB,,
+  754281,"01/03/2010, 06:27:59",HR,142.0,
+  754281,"06/20/2010, 20:23:50",BP//SYSTOLIC,134,
+  754281,"06/20/2010, 20:23:50",BP//DIASTOLIC,76,
+  754281,"06/20/2010, 21:00:02",TEMP//C,36.2,
+held_out/0: |-2
+  subject_id,time,code,numeric_value,text_value
+  1500733,,EYE_COLOR//BROWN,,
+  1500733,"07/20/1986, 00:00:00",DOB,,
+  1500733,"06/03/2010, 14:54:38",BP//SYSTOLIC,123,
+  1500733,"06/03/2010, 14:54:38",BP//DIASTOLIC,82,
+  1500733,"06/03/2010, 14:54:38",HR,91.4,
     """
 )
 
diff --git a/tests/MEDS_Transforms/test_multi_stage_preprocess_pipeline.py b/tests/MEDS_Transforms/test_multi_stage_preprocess_pipeline.py
index 6667313..7a1b1a7 100644
--- a/tests/MEDS_Transforms/test_multi_stage_preprocess_pipeline.py
+++ b/tests/MEDS_Transforms/test_multi_stage_preprocess_pipeline.py
@@ -81,150 +81,150 @@ fit_normalization:
 # After filtering out subjects with fewer than 5 events:
 WANT_FILTER = parse_shards_yaml(
     f"""
-  "filter_subjects/train/0": |-2
-    {subject_id_field},time,code,numeric_value
-    239684,,EYE_COLOR//BROWN,
-    239684,,HEIGHT,175.271115221764
-    239684,"12/28/1980, 00:00:00",DOB,
-    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,,EYE_COLOR//BLUE,
-    1195293,,HEIGHT,164.6868838269085
-    1195293,"06/20/1978, 00:00:00",DOB,
-    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,
+"filter_subjects/train/0": |-2
+  {subject_id_field},time,code,numeric_value
+  239684,,EYE_COLOR//BROWN,
+  239684,,HEIGHT,175.271115221764
+  239684,"12/28/1980, 00:00:00",DOB,
+  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,,EYE_COLOR//BLUE,
+  1195293,,HEIGHT,164.6868838269085
+  1195293,"06/20/1978, 00:00:00",DOB,
+  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,
 
-  "filter_subjects/train/1": |-2
-    {subject_id_field},time,code,numeric_value
+"filter_subjects/train/1": |-2
+  {subject_id_field},time,code,numeric_value
 
-  "filter_subjects/tuning/0": |-2
-    {subject_id_field},time,code,numeric_value
+"filter_subjects/tuning/0": |-2
+  {subject_id_field},time,code,numeric_value
 
-  "filter_subjects/held_out/0": |-2
-    {subject_id_field},time,code,numeric_value
-    1500733,,EYE_COLOR//BROWN,
-    1500733,,HEIGHT,158.60131573580904
-    1500733,"07/20/1986, 00:00:00",DOB,
-    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,
+"filter_subjects/held_out/0": |-2
+  {subject_id_field},time,code,numeric_value
+  1500733,,EYE_COLOR//BROWN,
+  1500733,,HEIGHT,158.60131573580904
+  1500733,"07/20/1986, 00:00:00",DOB,
+  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,
 """
 )
 
 WANT_TIME_DERIVED = parse_shards_yaml(
     f"""
-  "add_time_derived_measurements/train/0": |-2
-    {subject_id_field},time,code,numeric_value
-    239684,,EYE_COLOR//BROWN,
-    239684,,HEIGHT,175.271115221764
-    239684,"12/28/1980, 00:00:00","TIME_OF_DAY//[00,06)",
-    239684,"12/28/1980, 00:00:00",DOB,
-    239684,"05/11/2010, 17:41:51","TIME_OF_DAY//[12,18)",
-    239684,"05/11/2010, 17:41:51",AGE,29.36883360091833
-    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","TIME_OF_DAY//[12,18)",
-    239684,"05/11/2010, 17:48:48",AGE,29.36884681513314
-    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","TIME_OF_DAY//[18,24)",
-    239684,"05/11/2010, 18:25:35",AGE,29.36891675223647
-    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","TIME_OF_DAY//[18,24)",
-    239684,"05/11/2010, 18:57:18",AGE,29.36897705595538
-    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","TIME_OF_DAY//[18,24)",
-    239684,"05/11/2010, 19:27:19",AGE,29.369034127420306
-    239684,"05/11/2010, 19:27:19",DISCHARGE,
-    1195293,,EYE_COLOR//BLUE,
-    1195293,,HEIGHT,164.6868838269085
-    1195293,"06/20/1978, 00:00:00","TIME_OF_DAY//[00,06)",
-    1195293,"06/20/1978, 00:00:00",DOB,
-    1195293,"06/20/2010, 19:23:52","TIME_OF_DAY//[18,24)",
-    1195293,"06/20/2010, 19:23:52",AGE,32.002896271955265
-    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","TIME_OF_DAY//[18,24)",
-    1195293,"06/20/2010, 19:25:32",AGE,32.00289944083172
-    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","TIME_OF_DAY//[18,24)",
-    1195293,"06/20/2010, 19:45:19",AGE,32.00293705539522
-    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","TIME_OF_DAY//[18,24)",
-    1195293,"06/20/2010, 20:12:31",AGE,32.002988771458945
-    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","TIME_OF_DAY//[18,24)",
-    1195293,"06/20/2010, 20:24:44",AGE,32.00301199932335
-    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","TIME_OF_DAY//[18,24)",
-    1195293,"06/20/2010, 20:41:33",AGE,32.003043973286765
-    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","TIME_OF_DAY//[18,24)",
-    1195293,"06/20/2010, 20:50:04",AGE,32.00306016624544
-    1195293,"06/20/2010, 20:50:04",DISCHARGE,
+"add_time_derived_measurements/train/0": |-2
+  {subject_id_field},time,code,numeric_value
+  239684,,EYE_COLOR//BROWN,
+  239684,,HEIGHT,175.271115221764
+  239684,"12/28/1980, 00:00:00","TIME_OF_DAY//[00,06)",
+  239684,"12/28/1980, 00:00:00",DOB,
+  239684,"05/11/2010, 17:41:51","TIME_OF_DAY//[12,18)",
+  239684,"05/11/2010, 17:41:51",AGE,29.36883360091833
+  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","TIME_OF_DAY//[12,18)",
+  239684,"05/11/2010, 17:48:48",AGE,29.36884681513314
+  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","TIME_OF_DAY//[18,24)",
+  239684,"05/11/2010, 18:25:35",AGE,29.36891675223647
+  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","TIME_OF_DAY//[18,24)",
+  239684,"05/11/2010, 18:57:18",AGE,29.36897705595538
+  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","TIME_OF_DAY//[18,24)",
+  239684,"05/11/2010, 19:27:19",AGE,29.369034127420306
+  239684,"05/11/2010, 19:27:19",DISCHARGE,
+  1195293,,EYE_COLOR//BLUE,
+  1195293,,HEIGHT,164.6868838269085
+  1195293,"06/20/1978, 00:00:00","TIME_OF_DAY//[00,06)",
+  1195293,"06/20/1978, 00:00:00",DOB,
+  1195293,"06/20/2010, 19:23:52","TIME_OF_DAY//[18,24)",
+  1195293,"06/20/2010, 19:23:52",AGE,32.002896271955265
+  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","TIME_OF_DAY//[18,24)",
+  1195293,"06/20/2010, 19:25:32",AGE,32.00289944083172
+  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","TIME_OF_DAY//[18,24)",
+  1195293,"06/20/2010, 19:45:19",AGE,32.00293705539522
+  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","TIME_OF_DAY//[18,24)",
+  1195293,"06/20/2010, 20:12:31",AGE,32.002988771458945
+  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","TIME_OF_DAY//[18,24)",
+  1195293,"06/20/2010, 20:24:44",AGE,32.00301199932335
+  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","TIME_OF_DAY//[18,24)",
+  1195293,"06/20/2010, 20:41:33",AGE,32.003043973286765
+  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","TIME_OF_DAY//[18,24)",
+  1195293,"06/20/2010, 20:50:04",AGE,32.00306016624544
+  1195293,"06/20/2010, 20:50:04",DISCHARGE,
 
-  "add_time_derived_measurements/train/1": |-2
-    {subject_id_field},time,code,numeric_value
+"add_time_derived_measurements/train/1": |-2
+  {subject_id_field},time,code,numeric_value
 
-  "add_time_derived_measurements/tuning/0": |-2
-    {subject_id_field},time,code,numeric_value
+"add_time_derived_measurements/tuning/0": |-2
+  {subject_id_field},time,code,numeric_value
 
-  "add_time_derived_measurements/held_out/0": |-2
-    {subject_id_field},time,code,numeric_value
-    1500733,,EYE_COLOR//BROWN,
-    1500733,,HEIGHT,158.60131573580904
-    1500733,"07/20/1986, 00:00:00","TIME_OF_DAY//[00,06)",
-    1500733,"07/20/1986, 00:00:00",DOB,
-    1500733,"06/03/2010, 14:54:38","TIME_OF_DAY//[12,18)",
-    1500733,"06/03/2010, 14:54:38",AGE,23.873531791091356
-    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","TIME_OF_DAY//[12,18)",
-    1500733,"06/03/2010, 15:39:49",AGE,23.873617699332012
-    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","TIME_OF_DAY//[12,18)",
-    1500733,"06/03/2010, 16:20:49",AGE,23.873695653692767
-    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","TIME_OF_DAY//[12,18)",
-    1500733,"06/03/2010, 16:44:26",AGE,23.873740556672114
-    1500733,"06/03/2010, 16:44:26",DISCHARGE,
+"add_time_derived_measurements/held_out/0": |-2
+  {subject_id_field},time,code,numeric_value
+  1500733,,EYE_COLOR//BROWN,
+  1500733,,HEIGHT,158.60131573580904
+  1500733,"07/20/1986, 00:00:00","TIME_OF_DAY//[00,06)",
+  1500733,"07/20/1986, 00:00:00",DOB,
+  1500733,"06/03/2010, 14:54:38","TIME_OF_DAY//[12,18)",
+  1500733,"06/03/2010, 14:54:38",AGE,23.873531791091356
+  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","TIME_OF_DAY//[12,18)",
+  1500733,"06/03/2010, 15:39:49",AGE,23.873617699332012
+  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","TIME_OF_DAY//[12,18)",
+  1500733,"06/03/2010, 16:20:49",AGE,23.873695653692767
+  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","TIME_OF_DAY//[12,18)",
+  1500733,"06/03/2010, 16:44:26",AGE,23.873740556672114
+  1500733,"06/03/2010, 16:44:26",DISCHARGE,
 """
 )
 
@@ -389,93 +389,93 @@ shape: (4, 3)
 
 WANT_OCCLUDE_OUTLIERS = parse_shards_yaml(
     f"""
-  "occlude_outliers/train/0": |-2
-    {subject_id_field},time,code,numeric_value,numeric_value/is_inlier
-    239684,,EYE_COLOR//BROWN,,
-    239684,,HEIGHT,,false
-    239684,"12/28/1980, 00:00:00","TIME_OF_DAY//[00,06)",,
-    239684,"12/28/1980, 00:00:00",DOB,,
-    239684,"05/11/2010, 17:41:51","TIME_OF_DAY//[12,18)",,
-    239684,"05/11/2010, 17:41:51",AGE,,false
-    239684,"05/11/2010, 17:41:51",ADMISSION//CARDIAC,,
-    239684,"05/11/2010, 17:41:51",HR,,false
-    239684,"05/11/2010, 17:41:51",TEMP,,false
-    239684,"05/11/2010, 17:48:48","TIME_OF_DAY//[12,18)",,
-    239684,"05/11/2010, 17:48:48",AGE,,false
-    239684,"05/11/2010, 17:48:48",HR,,false
-    239684,"05/11/2010, 17:48:48",TEMP,,false
-    239684,"05/11/2010, 18:25:35","TIME_OF_DAY//[18,24)",,
-    239684,"05/11/2010, 18:25:35",AGE,,false
-    239684,"05/11/2010, 18:25:35",HR,113.4,true
-    239684,"05/11/2010, 18:25:35",TEMP,,false
-    239684,"05/11/2010, 18:57:18","TIME_OF_DAY//[18,24)",,
-    239684,"05/11/2010, 18:57:18",AGE,,false
-    239684,"05/11/2010, 18:57:18",HR,112.6,true
-    239684,"05/11/2010, 18:57:18",TEMP,,false
-    239684,"05/11/2010, 19:27:19","TIME_OF_DAY//[18,24)",,
-    239684,"05/11/2010, 19:27:19",AGE,,false
-    239684,"05/11/2010, 19:27:19",DISCHARGE,,
-    1195293,,EYE_COLOR//BLUE,,
-    1195293,,HEIGHT,,false
-    1195293,"06/20/1978, 00:00:00","TIME_OF_DAY//[00,06)",,
-    1195293,"06/20/1978, 00:00:00",DOB,,
-    1195293,"06/20/2010, 19:23:52","TIME_OF_DAY//[18,24)",,
-    1195293,"06/20/2010, 19:23:52",AGE,32.002896271955265,true
-    1195293,"06/20/2010, 19:23:52",ADMISSION//CARDIAC,,
-    1195293,"06/20/2010, 19:23:52",HR,109.0,true
-    1195293,"06/20/2010, 19:23:52",TEMP,100.0,true
-    1195293,"06/20/2010, 19:25:32","TIME_OF_DAY//[18,24)",,
-    1195293,"06/20/2010, 19:25:32",AGE,32.00289944083172,true
-    1195293,"06/20/2010, 19:25:32",HR,114.1,true
-    1195293,"06/20/2010, 19:25:32",TEMP,100.0,true
-    1195293,"06/20/2010, 19:45:19","TIME_OF_DAY//[18,24)",,
-    1195293,"06/20/2010, 19:45:19",AGE,32.00293705539522,true
-    1195293,"06/20/2010, 19:45:19",HR,,false
-    1195293,"06/20/2010, 19:45:19",TEMP,99.9,true
-    1195293,"06/20/2010, 20:12:31","TIME_OF_DAY//[18,24)",,
-    1195293,"06/20/2010, 20:12:31",AGE,32.002988771458945,true
-    1195293,"06/20/2010, 20:12:31",HR,112.5,true
-    1195293,"06/20/2010, 20:12:31",TEMP,99.8,true
-    1195293,"06/20/2010, 20:24:44","TIME_OF_DAY//[18,24)",
-    1195293,"06/20/2010, 20:24:44",AGE,32.00301199932335,true
-    1195293,"06/20/2010, 20:24:44",HR,107.7,true
-    1195293,"06/20/2010, 20:24:44",TEMP,100.0,true
-    1195293,"06/20/2010, 20:41:33","TIME_OF_DAY//[18,24)",,
-    1195293,"06/20/2010, 20:41:33",AGE,32.003043973286765,true
-    1195293,"06/20/2010, 20:41:33",HR,107.5,true
-    1195293,"06/20/2010, 20:41:33",TEMP,100.4,true
-    1195293,"06/20/2010, 20:50:04","TIME_OF_DAY//[18,24)",,
-    1195293,"06/20/2010, 20:50:04",AGE,32.00306016624544,true
-    1195293,"06/20/2010, 20:50:04",DISCHARGE,,
+"occlude_outliers/train/0": |-2
+  {subject_id_field},time,code,numeric_value,numeric_value/is_inlier
+  239684,,EYE_COLOR//BROWN,,
+  239684,,HEIGHT,,false
+  239684,"12/28/1980, 00:00:00","TIME_OF_DAY//[00,06)",,
+  239684,"12/28/1980, 00:00:00",DOB,,
+  239684,"05/11/2010, 17:41:51","TIME_OF_DAY//[12,18)",,
+  239684,"05/11/2010, 17:41:51",AGE,,false
+  239684,"05/11/2010, 17:41:51",ADMISSION//CARDIAC,,
+  239684,"05/11/2010, 17:41:51",HR,,false
+  239684,"05/11/2010, 17:41:51",TEMP,,false
+  239684,"05/11/2010, 17:48:48","TIME_OF_DAY//[12,18)",,
+  239684,"05/11/2010, 17:48:48",AGE,,false
+  239684,"05/11/2010, 17:48:48",HR,,false
+  239684,"05/11/2010, 17:48:48",TEMP,,false
+  239684,"05/11/2010, 18:25:35","TIME_OF_DAY//[18,24)",,
+  239684,"05/11/2010, 18:25:35",AGE,,false
+  239684,"05/11/2010, 18:25:35",HR,113.4,true
+  239684,"05/11/2010, 18:25:35",TEMP,,false
+  239684,"05/11/2010, 18:57:18","TIME_OF_DAY//[18,24)",,
+  239684,"05/11/2010, 18:57:18",AGE,,false
+  239684,"05/11/2010, 18:57:18",HR,112.6,true
+  239684,"05/11/2010, 18:57:18",TEMP,,false
+  239684,"05/11/2010, 19:27:19","TIME_OF_DAY//[18,24)",,
+  239684,"05/11/2010, 19:27:19",AGE,,false
+  239684,"05/11/2010, 19:27:19",DISCHARGE,,
+  1195293,,EYE_COLOR//BLUE,,
+  1195293,,HEIGHT,,false
+  1195293,"06/20/1978, 00:00:00","TIME_OF_DAY//[00,06)",,
+  1195293,"06/20/1978, 00:00:00",DOB,,
+  1195293,"06/20/2010, 19:23:52","TIME_OF_DAY//[18,24)",,
+  1195293,"06/20/2010, 19:23:52",AGE,32.002896271955265,true
+  1195293,"06/20/2010, 19:23:52",ADMISSION//CARDIAC,,
+  1195293,"06/20/2010, 19:23:52",HR,109.0,true
+  1195293,"06/20/2010, 19:23:52",TEMP,100.0,true
+  1195293,"06/20/2010, 19:25:32","TIME_OF_DAY//[18,24)",,
+  1195293,"06/20/2010, 19:25:32",AGE,32.00289944083172,true
+  1195293,"06/20/2010, 19:25:32",HR,114.1,true
+  1195293,"06/20/2010, 19:25:32",TEMP,100.0,true
+  1195293,"06/20/2010, 19:45:19","TIME_OF_DAY//[18,24)",,
+  1195293,"06/20/2010, 19:45:19",AGE,32.00293705539522,true
+  1195293,"06/20/2010, 19:45:19",HR,,false
+  1195293,"06/20/2010, 19:45:19",TEMP,99.9,true
+  1195293,"06/20/2010, 20:12:31","TIME_OF_DAY//[18,24)",,
+  1195293,"06/20/2010, 20:12:31",AGE,32.002988771458945,true
+  1195293,"06/20/2010, 20:12:31",HR,112.5,true
+  1195293,"06/20/2010, 20:12:31",TEMP,99.8,true
+  1195293,"06/20/2010, 20:24:44","TIME_OF_DAY//[18,24)",
+  1195293,"06/20/2010, 20:24:44",AGE,32.00301199932335,true
+  1195293,"06/20/2010, 20:24:44",HR,107.7,true
+  1195293,"06/20/2010, 20:24:44",TEMP,100.0,true
+  1195293,"06/20/2010, 20:41:33","TIME_OF_DAY//[18,24)",,
+  1195293,"06/20/2010, 20:41:33",AGE,32.003043973286765,true
+  1195293,"06/20/2010, 20:41:33",HR,107.5,true
+  1195293,"06/20/2010, 20:41:33",TEMP,100.4,true
+  1195293,"06/20/2010, 20:50:04","TIME_OF_DAY//[18,24)",,
+  1195293,"06/20/2010, 20:50:04",AGE,32.00306016624544,true
+  1195293,"06/20/2010, 20:50:04",DISCHARGE,,
 
-  "occlude_outliers/train/1": |-2
-    {subject_id_field},time,code,numeric_value,numeric_value/is_inlier
+"occlude_outliers/train/1": |-2
+  {subject_id_field},time,code,numeric_value,numeric_value/is_inlier
 
-  "occlude_outliers/tuning/0": |-2
-    {subject_id_field},time,code,numeric_value,numeric_value/is_inlier
+"occlude_outliers/tuning/0": |-2
+  {subject_id_field},time,code,numeric_value,numeric_value/is_inlier
 
-  "occlude_outliers/held_out/0": |-2
-    {subject_id_field},time,code,numeric_value,numeric_value/is_inlier
-    1500733,,EYE_COLOR//BROWN,,
-    1500733,,HEIGHT,,false
-    1500733,"07/20/1986, 00:00:00","TIME_OF_DAY//[00,06)",,
-    1500733,"07/20/1986, 00:00:00",DOB,,
-    1500733,"06/03/2010, 14:54:38","TIME_OF_DAY//[12,18)",,
-    1500733,"06/03/2010, 14:54:38",AGE,,false
-    1500733,"06/03/2010, 14:54:38",ADMISSION//ORTHOPEDIC,,
-    1500733,"06/03/2010, 14:54:38",HR,,false
-    1500733,"06/03/2010, 14:54:38",TEMP,100.0,true
-    1500733,"06/03/2010, 15:39:49","TIME_OF_DAY//[12,18)",,
-    1500733,"06/03/2010, 15:39:49",AGE,,false
-    1500733,"06/03/2010, 15:39:49",HR,,false
-    1500733,"06/03/2010, 15:39:49",TEMP,100.3,true
-    1500733,"06/03/2010, 16:20:49","TIME_OF_DAY//[12,18)",,
-    1500733,"06/03/2010, 16:20:49",AGE,,false
-    1500733,"06/03/2010, 16:20:49",HR,,false
-    1500733,"06/03/2010, 16:20:49",TEMP,100.1,true
-    1500733,"06/03/2010, 16:44:26","TIME_OF_DAY//[12,18)",,
-    1500733,"06/03/2010, 16:44:26",AGE,,false
-    1500733,"06/03/2010, 16:44:26",DISCHARGE,,
+"occlude_outliers/held_out/0": |-2
+  {subject_id_field},time,code,numeric_value,numeric_value/is_inlier
+  1500733,,EYE_COLOR//BROWN,,
+  1500733,,HEIGHT,,false
+  1500733,"07/20/1986, 00:00:00","TIME_OF_DAY//[00,06)",,
+  1500733,"07/20/1986, 00:00:00",DOB,,
+  1500733,"06/03/2010, 14:54:38","TIME_OF_DAY//[12,18)",,
+  1500733,"06/03/2010, 14:54:38",AGE,,false
+  1500733,"06/03/2010, 14:54:38",ADMISSION//ORTHOPEDIC,,
+  1500733,"06/03/2010, 14:54:38",HR,,false
+  1500733,"06/03/2010, 14:54:38",TEMP,100.0,true
+  1500733,"06/03/2010, 15:39:49","TIME_OF_DAY//[12,18)",,
+  1500733,"06/03/2010, 15:39:49",AGE,,false
+  1500733,"06/03/2010, 15:39:49",HR,,false
+  1500733,"06/03/2010, 15:39:49",TEMP,100.3,true
+  1500733,"06/03/2010, 16:20:49","TIME_OF_DAY//[12,18)",,
+  1500733,"06/03/2010, 16:20:49",AGE,,false
+  1500733,"06/03/2010, 16:20:49",HR,,false
+  1500733,"06/03/2010, 16:20:49",TEMP,100.1,true
+  1500733,"06/03/2010, 16:44:26","TIME_OF_DAY//[12,18)",,
+  1500733,"06/03/2010, 16:44:26",AGE,,false
+  1500733,"06/03/2010, 16:44:26",DISCHARGE,,
 """
 )
 
@@ -774,93 +774,93 @@ held_out/0
 # Note we have dropped the row in the held out shard that doesn't have a code in the vocabulary!
 WANT_NORMALIZATION = parse_shards_yaml(
     f"""
-  "normalization/train/0": |-2
-    {subject_id_field},time,code,numeric_value
-    239684,,6,
-    239684,,7,
-    239684,"12/28/1980, 00:00:00",10,
-    239684,"12/28/1980, 00:00:00",4,
-    239684,"05/11/2010, 17:41:51",11,
-    239684,"05/11/2010, 17:41:51",2,
-    239684,"05/11/2010, 17:41:51",1,
-    239684,"05/11/2010, 17:41:51",8,
-    239684,"05/11/2010, 17:41:51",9,
-    239684,"05/11/2010, 17:48:48",11,
-    239684,"05/11/2010, 17:48:48",2,
-    239684,"05/11/2010, 17:48:48",8,
-    239684,"05/11/2010, 17:48:48",9,
-    239684,"05/11/2010, 18:25:35",12,
-    239684,"05/11/2010, 18:25:35",2,
-    239684,"05/11/2010, 18:25:35",8,0.9341503977775574
-    239684,"05/11/2010, 18:25:35",9,
-    239684,"05/11/2010, 18:57:18",12,
-    239684,"05/11/2010, 18:57:18",2,
-    239684,"05/11/2010, 18:57:18",8,0.6264293789863586
-    239684,"05/11/2010, 18:57:18",9,
-    239684,"05/11/2010, 19:27:19",12,
-    239684,"05/11/2010, 19:27:19",2,
-    239684,"05/11/2010, 19:27:19",3,
-    1195293,,5,
-    1195293,,7,
-    1195293,"06/20/1978, 00:00:00",10,
-    1195293,"06/20/1978, 00:00:00",4,
-    1195293,"06/20/2010, 19:23:52",12,
-    1195293,"06/20/2010, 19:23:52",2,nan
-    1195293,"06/20/2010, 19:23:52",1,
-    1195293,"06/20/2010, 19:23:52",8,-0.7583094239234924
-    1195293,"06/20/2010, 19:23:52",9,-0.0889078751206398
-    1195293,"06/20/2010, 19:25:32",12,
-    1195293,"06/20/2010, 19:25:32",2,nan
-    1195293,"06/20/2010, 19:25:32",8,1.2034040689468384
-    1195293,"06/20/2010, 19:25:32",9,-0.0889078751206398
-    1195293,"06/20/2010, 19:45:19",12,
-    1195293,"06/20/2010, 19:45:19",2,nan
-    1195293,"06/20/2010, 19:45:19",8,
-    1195293,"06/20/2010, 19:45:19",9,-0.6222330927848816
-    1195293,"06/20/2010, 20:12:31",12,
-    1195293,"06/20/2010, 20:12:31",2,nan
-    1195293,"06/20/2010, 20:12:31",8,0.5879650115966797
-    1195293,"06/20/2010, 20:12:31",9,-1.1555582284927368
-    1195293,"06/20/2010, 20:24:44",12
-    1195293,"06/20/2010, 20:24:44",2,nan
-    1195293,"06/20/2010, 20:24:44",8,-1.2583553791046143
-    1195293,"06/20/2010, 20:24:44",9,-0.0889078751206398
-    1195293,"06/20/2010, 20:41:33",12,
-    1195293,"06/20/2010, 20:41:33",2,nan
-    1195293,"06/20/2010, 20:41:33",8,-1.3352841138839722
-    1195293,"06/20/2010, 20:41:33",9,2.04443359375
-    1195293,"06/20/2010, 20:50:04",12,
-    1195293,"06/20/2010, 20:50:04",2,nan
-    1195293,"06/20/2010, 20:50:04",3,
+"normalization/train/0": |-2
+  {subject_id_field},time,code,numeric_value
+  239684,,6,
+  239684,,7,
+  239684,"12/28/1980, 00:00:00",10,
+  239684,"12/28/1980, 00:00:00",4,
+  239684,"05/11/2010, 17:41:51",11,
+  239684,"05/11/2010, 17:41:51",2,
+  239684,"05/11/2010, 17:41:51",1,
+  239684,"05/11/2010, 17:41:51",8,
+  239684,"05/11/2010, 17:41:51",9,
+  239684,"05/11/2010, 17:48:48",11,
+  239684,"05/11/2010, 17:48:48",2,
+  239684,"05/11/2010, 17:48:48",8,
+  239684,"05/11/2010, 17:48:48",9,
+  239684,"05/11/2010, 18:25:35",12,
+  239684,"05/11/2010, 18:25:35",2,
+  239684,"05/11/2010, 18:25:35",8,0.9341503977775574
+  239684,"05/11/2010, 18:25:35",9,
+  239684,"05/11/2010, 18:57:18",12,
+  239684,"05/11/2010, 18:57:18",2,
+  239684,"05/11/2010, 18:57:18",8,0.6264293789863586
+  239684,"05/11/2010, 18:57:18",9,
+  239684,"05/11/2010, 19:27:19",12,
+  239684,"05/11/2010, 19:27:19",2,
+  239684,"05/11/2010, 19:27:19",3,
+  1195293,,5,
+  1195293,,7,
+  1195293,"06/20/1978, 00:00:00",10,
+  1195293,"06/20/1978, 00:00:00",4,
+  1195293,"06/20/2010, 19:23:52",12,
+  1195293,"06/20/2010, 19:23:52",2,nan
+  1195293,"06/20/2010, 19:23:52",1,
+  1195293,"06/20/2010, 19:23:52",8,-0.7583094239234924
+  1195293,"06/20/2010, 19:23:52",9,-0.0889078751206398
+  1195293,"06/20/2010, 19:25:32",12,
+  1195293,"06/20/2010, 19:25:32",2,nan
+  1195293,"06/20/2010, 19:25:32",8,1.2034040689468384
+  1195293,"06/20/2010, 19:25:32",9,-0.0889078751206398
+  1195293,"06/20/2010, 19:45:19",12,
+  1195293,"06/20/2010, 19:45:19",2,nan
+  1195293,"06/20/2010, 19:45:19",8,
+  1195293,"06/20/2010, 19:45:19",9,-0.6222330927848816
+  1195293,"06/20/2010, 20:12:31",12,
+  1195293,"06/20/2010, 20:12:31",2,nan
+  1195293,"06/20/2010, 20:12:31",8,0.5879650115966797
+  1195293,"06/20/2010, 20:12:31",9,-1.1555582284927368
+  1195293,"06/20/2010, 20:24:44",12
+  1195293,"06/20/2010, 20:24:44",2,nan
+  1195293,"06/20/2010, 20:24:44",8,-1.2583553791046143
+  1195293,"06/20/2010, 20:24:44",9,-0.0889078751206398
+  1195293,"06/20/2010, 20:41:33",12,
+  1195293,"06/20/2010, 20:41:33",2,nan
+  1195293,"06/20/2010, 20:41:33",8,-1.3352841138839722
+  1195293,"06/20/2010, 20:41:33",9,2.04443359375
+  1195293,"06/20/2010, 20:50:04",12,
+  1195293,"06/20/2010, 20:50:04",2,nan
+  1195293,"06/20/2010, 20:50:04",3,
 
-  "normalization/train/1": |-2
-    {subject_id_field},time,code,numeric_value
+"normalization/train/1": |-2
+  {subject_id_field},time,code,numeric_value
 
-  "normalization/tuning/0": |-2
-    {subject_id_field},time,code,numeric_value
+"normalization/tuning/0": |-2
+  {subject_id_field},time,code,numeric_value
 
-  "normalization/held_out/0": |-2
-    {subject_id_field},time,code,numeric_value
-    1500733,,6,
-    1500733,,7,
-    1500733,"07/20/1986, 00:00:00",10,
-    1500733,"07/20/1986, 00:00:00",4,
-    1500733,"06/03/2010, 14:54:38",11,
-    1500733,"06/03/2010, 14:54:38",2,
-    1500733,"06/03/2010, 14:54:38",8,
-    1500733,"06/03/2010, 14:54:38",9,-0.0889078751206398
-    1500733,"06/03/2010, 15:39:49",11,
-    1500733,"06/03/2010, 15:39:49",2,
-    1500733,"06/03/2010, 15:39:49",8,
-    1500733,"06/03/2010, 15:39:49",9,1.5111083984375
-    1500733,"06/03/2010, 16:20:49",11,
-    1500733,"06/03/2010, 16:20:49",2,
-    1500733,"06/03/2010, 16:20:49",8,
-    1500733,"06/03/2010, 16:20:49",9,0.4444173276424408
-    1500733,"06/03/2010, 16:44:26",11,
-    1500733,"06/03/2010, 16:44:26",2,
-    1500733,"06/03/2010, 16:44:26",3,
-    """,
+"normalization/held_out/0": |-2
+  {subject_id_field},time,code,numeric_value
+  1500733,,6,
+  1500733,,7,
+  1500733,"07/20/1986, 00:00:00",10,
+  1500733,"07/20/1986, 00:00:00",4,
+  1500733,"06/03/2010, 14:54:38",11,
+  1500733,"06/03/2010, 14:54:38",2,
+  1500733,"06/03/2010, 14:54:38",8,
+  1500733,"06/03/2010, 14:54:38",9,-0.0889078751206398
+  1500733,"06/03/2010, 15:39:49",11,
+  1500733,"06/03/2010, 15:39:49",2,
+  1500733,"06/03/2010, 15:39:49",8,
+  1500733,"06/03/2010, 15:39:49",9,1.5111083984375
+  1500733,"06/03/2010, 16:20:49",11,
+  1500733,"06/03/2010, 16:20:49",2,
+  1500733,"06/03/2010, 16:20:49",8,
+  1500733,"06/03/2010, 16:20:49",9,0.4444173276424408
+  1500733,"06/03/2010, 16:44:26",11,
+  1500733,"06/03/2010, 16:44:26",2,
+  1500733,"06/03/2010, 16:44:26",3,
+  """,
     code=pl.UInt8,
 )
 
diff --git a/tests/MEDS_Transforms/transform_tester_base.py b/tests/MEDS_Transforms/transform_tester_base.py
index 6a1d4ab..93e70aa 100644
--- a/tests/MEDS_Transforms/transform_tester_base.py
+++ b/tests/MEDS_Transforms/transform_tester_base.py
@@ -17,10 +17,7 @@ from pathlib import Path
 import polars as pl
 from meds import subject_id_field
 
-from tests.utils import FILE_T, multi_stage_tester, parse_meds_csvs, parse_shards_yaml, single_stage_tester
-
-# So it can be imported from here
-parse_shards_yaml = parse_shards_yaml
+from tests.utils import FILE_T, multi_stage_tester, parse_shards_yaml, single_stage_tester
 
 # Test MEDS data (inputs)
 
@@ -31,99 +28,90 @@ SHARDS = {
     "held_out/0": [1500733],
 }
 
-SPLITS = {
-    "train": [239684, 1195293, 68729, 814703],
-    "tuning": [754281],
-    "held_out": [1500733],
-}
-
-MEDS_TRAIN_0 = """
-subject_id,time,code,numeric_value
-239684,,EYE_COLOR//BROWN,
-239684,,HEIGHT,175.271115221764
-239684,"12/28/1980, 00:00:00",DOB,
-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,,EYE_COLOR//BLUE,
-1195293,,HEIGHT,164.6868838269085
-1195293,"06/20/1978, 00:00:00",DOB,
-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_TRAIN_1 = """
-subject_id,time,code,numeric_value
-68729,,EYE_COLOR//HAZEL,
-68729,,HEIGHT,160.3953106166676
-68729,"03/09/1978, 00:00:00",DOB,
-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,,EYE_COLOR//HAZEL,
-814703,,HEIGHT,156.48559093209357
-814703,"03/28/1976, 00:00:00",DOB,
-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_TUNING_0 = """
-subject_id,time,code,numeric_value
-754281,,EYE_COLOR//BROWN,
-754281,,HEIGHT,166.22261567137025
-754281,"12/19/1988, 00:00:00",DOB,
-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_HELD_OUT_0 = """
-subject_id,time,code,numeric_value
-1500733,,EYE_COLOR//BROWN,
-1500733,,HEIGHT,158.60131573580904
-1500733,"07/20/1986, 00:00:00",DOB,
-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_SHARDS = parse_meds_csvs(
+SPLITS_DF = pl.DataFrame(
     {
-        "train/0": MEDS_TRAIN_0,
-        "train/1": MEDS_TRAIN_1,
-        "tuning/0": MEDS_TUNING_0,
-        "held_out/0": MEDS_HELD_OUT_0,
+        subject_id_field: [239684, 1195293, 68729, 814703, 754281, 1500733],
+        "split": ["train", "train", "train", "train", "tuning", "held_out"],
     }
 )
 
+MEDS_SHARDS = parse_shards_yaml(
+    """
+train/0: |-2
+  subject_id,time,code,numeric_value
+  239684,,EYE_COLOR//BROWN,
+  239684,,HEIGHT,175.271115221764
+  239684,"12/28/1980, 00:00:00",DOB,
+  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,,EYE_COLOR//BLUE,
+  1195293,,HEIGHT,164.6868838269085
+  1195293,"06/20/1978, 00:00:00",DOB,
+  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,
+
+train/1: |-2
+  subject_id,time,code,numeric_value
+  68729,,EYE_COLOR//HAZEL,
+  68729,,HEIGHT,160.3953106166676
+  68729,"03/09/1978, 00:00:00",DOB,
+  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,,EYE_COLOR//HAZEL,
+  814703,,HEIGHT,156.48559093209357
+  814703,"03/28/1976, 00:00:00",DOB,
+  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,
+
+tuning/0: |-2
+  subject_id,time,code,numeric_value
+  754281,,EYE_COLOR//BROWN,
+  754281,,HEIGHT,166.22261567137025
+  754281,"12/19/1988, 00:00:00",DOB,
+  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,
+
+held_out/0: |-2
+  subject_id,time,code,numeric_value
+  1500733,,EYE_COLOR//BROWN,
+  1500733,,HEIGHT,158.60131573580904
+  1500733,"07/20/1986, 00:00:00",DOB,
+  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_CODE_METADATA_CSV = """
 code,code/n_occurrences,code/n_subjects,values/n_occurrences,values/sum,values/sum_sqd,description,parent_codes
@@ -197,14 +185,17 @@ def remap_inputs_for_transform(
     unified_inputs["metadata/.shards.json"] = input_shards_map
 
     if input_splits_map is None:
-        input_splits_map = SPLITS
+        input_splits_map = SPLITS_DF
 
-    input_splits_as_df = defaultdict(list)
-    for split_name, subject_ids in input_splits_map.items():
-        input_splits_as_df[subject_id_field].extend(subject_ids)
-        input_splits_as_df["split"].extend([split_name] * len(subject_ids))
+    if isinstance(input_splits_map, pl.DataFrame):
+        input_splits_df = input_splits_map
+    else:
+        input_splits_as_df = defaultdict(list)
+        for split_name, subject_ids in input_splits_map.items():
+            input_splits_as_df[subject_id_field].extend(subject_ids)
+            input_splits_as_df["split"].extend([split_name] * len(subject_ids))
 
-    input_splits_df = pl.DataFrame(input_splits_as_df)
+        input_splits_df = pl.DataFrame(input_splits_as_df)
 
     unified_inputs["metadata/subject_splits.parquet"] = input_splits_df
 
diff --git a/tests/__init__.py b/tests/__init__.py
index e69de29..8952037 100644
--- a/tests/__init__.py
+++ b/tests/__init__.py
@@ -0,0 +1,14 @@
+import os
+
+import rootutils
+
+root = rootutils.setup_root(__file__, dotenv=True, pythonpath=True, cwd=True)
+
+code_root = root / "src" / "MEDS_transforms"
+
+USE_LOCAL_SCRIPTS = os.environ.get("DO_USE_LOCAL_SCRIPTS", "0") == "1"
+
+if USE_LOCAL_SCRIPTS:
+    RUNNER_SCRIPT = code_root / "runner.py"
+else:
+    RUNNER_SCRIPT = "MEDS_transform-runner"
diff --git a/tests/test_with_runner.py b/tests/test_with_runner.py
new file mode 100644
index 0000000..5e2a4a1
--- /dev/null
+++ b/tests/test_with_runner.py
@@ -0,0 +1,273 @@
+"""Tests a multi-stage pre-processing pipeline via the Runner utility. Only checks final outputs.
+
+Set the bash env variable `DO_USE_LOCAL_SCRIPTS=1` to use the local py files, rather than the installed
+scripts.
+
+In this test, the following stages are run:
+  - filter_subjects
+  - add_time_derived_measurements
+  - fit_outlier_detection
+  - occlude_outliers
+  - fit_normalization
+  - fit_vocabulary_indices
+  - normalization
+  - tokenization
+  - tensorization
+
+The stage configuration arguments will be as given in the yaml block below:
+"""
+
+
+from functools import partial
+
+from tests import RUNNER_SCRIPT, USE_LOCAL_SCRIPTS
+from tests.MEDS_Transforms import (
+    ADD_TIME_DERIVED_MEASUREMENTS_SCRIPT,
+    AGGREGATE_CODE_METADATA_SCRIPT,
+    FILTER_SUBJECTS_SCRIPT,
+    FIT_VOCABULARY_INDICES_SCRIPT,
+    NORMALIZATION_SCRIPT,
+    OCCLUDE_OUTLIERS_SCRIPT,
+    TOKENIZATION_SCRIPT,
+)
+from tests.MEDS_Transforms.test_multi_stage_preprocess_pipeline import (
+    MEDS_CODE_METADATA,
+    WANT_FILTER,
+    WANT_FIT_NORMALIZATION,
+    WANT_FIT_OUTLIERS,
+    WANT_FIT_VOCABULARY_INDICES,
+    WANT_NORMALIZATION,
+    WANT_OCCLUDE_OUTLIERS,
+    WANT_TIME_DERIVED,
+    WANT_TOKENIZATION_EVENT_SEQS,
+    WANT_TOKENIZATION_SCHEMAS,
+    WANT_NRTs,
+)
+from tests.MEDS_Transforms.transform_tester_base import MEDS_SHARDS, SPLITS_DF
+from tests.utils import add_params, exact_str_regex, single_stage_tester
+
+# Normally, you wouldn't need to specify all of these scripts, but in testing with local scripts we need to
+# specify them all as they need to point to their python paths.
+if USE_LOCAL_SCRIPTS:
+    STAGE_RUNNER_YAML = f"""
+filter_subjects:
+  script: "python {FILTER_SUBJECTS_SCRIPT}"
+
+add_time_derived_measurements:
+  script: "python {ADD_TIME_DERIVED_MEASUREMENTS_SCRIPT}"
+
+occlude_outliers:
+  script: "python {OCCLUDE_OUTLIERS_SCRIPT}"
+
+fit_normalization:
+  script: "python {AGGREGATE_CODE_METADATA_SCRIPT}"
+
+fit_vocabulary_indices:
+  script: "python {FIT_VOCABULARY_INDICES_SCRIPT}"
+
+normalization:
+  script: "python {NORMALIZATION_SCRIPT}"
+
+tokenization:
+  script: "python {TOKENIZATION_SCRIPT}"
+    """
+else:
+    STAGE_RUNNER_YAML = f"""
+fit_normalization:
+  script: {AGGREGATE_CODE_METADATA_SCRIPT}
+    """
+
+PARALLEL_STAGE_RUNNER_YAML = f"""
+parallelize:
+  n_workers: 2
+  launcher: "joblib"
+
+{STAGE_RUNNER_YAML}
+"""
+
+
+PIPELINE_YAML = f"""
+defaults:
+  - _preprocess
+  - _self_
+
+input_dir: {{input_dir}}
+cohort_dir: {{cohort_dir}}
+
+description: "A test pipeline for the MEDS-transforms pipeline runner."
+
+stages:
+  - filter_subjects
+  - add_time_derived_measurements
+  - fit_outlier_detection
+  - occlude_outliers
+  - fit_normalization
+  - fit_vocabulary_indices
+  - normalization
+  - tokenization
+  - tensorization
+
+stage_configs:
+  filter_subjects:
+    min_events_per_subject: 5
+  add_time_derived_measurements:
+    age:
+      DOB_code: "DOB" # This is the MEDS official code for BIRTH
+      age_code: "AGE"
+      age_unit: "years"
+    time_of_day:
+      time_of_day_code: "TIME_OF_DAY"
+      endpoints: [6, 12, 18, 24]
+  fit_outlier_detection:
+    _script: {("python " if USE_LOCAL_SCRIPTS else "") + str(AGGREGATE_CODE_METADATA_SCRIPT)}
+    aggregations:
+      - "values/n_occurrences"
+      - "values/sum"
+      - "values/sum_sqd"
+  occlude_outliers:
+    stddev_cutoff: 1
+  fit_normalization:
+    aggregations:
+      - "code/n_occurrences"
+      - "code/n_subjects"
+      - "values/n_occurrences"
+      - "values/sum"
+      - "values/sum_sqd"
+"""
+
+NO_ARGS_HELP_STR = """
+== MEDS-Transforms Pipeline Runner ==
+MEDS-Transforms Pipeline Runner is a command line tool for running entire MEDS-transform pipelines in a single
+command.
+
+Runs the entire pipeline, end-to-end, based on the configuration provided.
+
+This script will launch many subsidiary commands via `subprocess`, one for each stage of the specified
+pipeline.
+
+**MEDS-transforms Pipeline description:**
+
+No description provided.
+"""
+
+WITH_CONFIG_HELP_STR = """
+== MEDS-Transforms Pipeline Runner ==
+MEDS-Transforms Pipeline Runner is a command line tool for running entire MEDS-transform pipelines in a single
+command.
+
+Runs the entire pipeline, end-to-end, based on the configuration provided.
+
+This script will launch many subsidiary commands via `subprocess`, one for each stage of the specified
+pipeline.
+
+**MEDS-transforms Pipeline description:**
+
+A test pipeline for the MEDS-transforms pipeline runner.
+"""
+
+
+def test_pipeline():
+    single_stage_tester(
+        script=str(RUNNER_SCRIPT) + " -h",
+        config_name="runner",
+        stage_name=None,
+        stage_kwargs=None,
+        do_pass_stage_name=False,
+        do_use_config_yaml=False,
+        input_files={},
+        want_outputs={},
+        assert_no_other_outputs=True,
+        should_error=False,
+        test_name="Runner Help Test",
+        do_include_dirs=False,
+        hydra_verbose=False,
+        stdout_regex=exact_str_regex(NO_ARGS_HELP_STR.strip()),
+    )
+
+    single_stage_tester(
+        script=str(RUNNER_SCRIPT) + " -h",
+        config_name="runner",
+        stage_name=None,
+        stage_kwargs=None,
+        do_pass_stage_name=False,
+        do_use_config_yaml=False,
+        input_files={"pipeline.yaml": partial(add_params, PIPELINE_YAML)},
+        want_outputs={},
+        assert_no_other_outputs=True,
+        should_error=False,
+        pipeline_config_fp="{input_dir}/pipeline.yaml",
+        test_name="Runner Help Test",
+        do_include_dirs=False,
+        hydra_verbose=False,
+        stdout_regex=exact_str_regex(WITH_CONFIG_HELP_STR.strip()),
+    )
+
+    single_stage_tester(
+        script=RUNNER_SCRIPT,
+        config_name="runner",
+        stage_name=None,
+        stage_kwargs=None,
+        do_pass_stage_name=False,
+        do_use_config_yaml=False,
+        input_files={
+            **{f"data/{k}": v for k, v in MEDS_SHARDS.items()},
+            "metadata/codes.parquet": MEDS_CODE_METADATA,
+            "metadata/subject_splits.parquet": SPLITS_DF,
+            "pipeline.yaml": partial(add_params, PIPELINE_YAML),
+            "stage_runner.yaml": STAGE_RUNNER_YAML,
+        },
+        want_outputs={
+            **WANT_FIT_NORMALIZATION,
+            **WANT_FIT_OUTLIERS,
+            **WANT_FIT_VOCABULARY_INDICES,
+            **WANT_FILTER,
+            **WANT_TIME_DERIVED,
+            **WANT_OCCLUDE_OUTLIERS,
+            **WANT_NORMALIZATION,
+            **WANT_TOKENIZATION_SCHEMAS,
+            **WANT_TOKENIZATION_EVENT_SEQS,
+            **WANT_NRTs,
+        },
+        assert_no_other_outputs=False,
+        should_error=False,
+        pipeline_config_fp="{input_dir}/pipeline.yaml",
+        stage_runner_fp="{input_dir}/stage_runner.yaml",
+        test_name="Runner Test",
+        do_include_dirs=False,
+        df_check_kwargs={"check_column_order": False},
+    )
+
+    single_stage_tester(
+        script=RUNNER_SCRIPT,
+        config_name="runner",
+        stage_name=None,
+        stage_kwargs=None,
+        do_pass_stage_name=False,
+        do_use_config_yaml=False,
+        input_files={
+            **{f"data/{k}": v for k, v in MEDS_SHARDS.items()},
+            "metadata/codes.parquet": MEDS_CODE_METADATA,
+            "metadata/subject_splits.parquet": SPLITS_DF,
+            "pipeline.yaml": partial(add_params, PIPELINE_YAML),
+            "stage_runner.yaml": PARALLEL_STAGE_RUNNER_YAML,
+        },
+        want_outputs={
+            **WANT_FIT_NORMALIZATION,
+            **WANT_FIT_OUTLIERS,
+            **WANT_FIT_VOCABULARY_INDICES,
+            **WANT_FILTER,
+            **WANT_TIME_DERIVED,
+            **WANT_OCCLUDE_OUTLIERS,
+            **WANT_NORMALIZATION,
+            **WANT_TOKENIZATION_SCHEMAS,
+            **WANT_TOKENIZATION_EVENT_SEQS,
+            **WANT_NRTs,
+        },
+        assert_no_other_outputs=False,
+        should_error=False,
+        pipeline_config_fp="{input_dir}/pipeline.yaml",
+        stage_runner_fp="{input_dir}/stage_runner.yaml",
+        test_name="Runner Test with parallelism",
+        do_include_dirs=False,
+        df_check_kwargs={"check_column_order": False},
+    )
diff --git a/tests/utils.py b/tests/utils.py
index 450a86b..a5018b4 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -1,4 +1,5 @@
 import json
+import re
 import subprocess
 import tempfile
 from contextlib import contextmanager
@@ -31,6 +32,10 @@ MEDS_PL_SCHEMA = {
 }
 
 
+def exact_str_regex(s: str) -> str:
+    return f"^{re.escape(s)}$"
+
+
 def parse_meds_csvs(
     csvs: str | dict[str, str], schema: dict[str, pl.DataType] = MEDS_PL_SCHEMA
 ) -> pl.DataFrame | dict[str, pl.DataFrame]:
@@ -57,7 +62,7 @@ def parse_meds_csvs(
 
 def parse_shards_yaml(yaml_str: str, **schema_updates) -> pl.DataFrame:
     schema = {**MEDS_PL_SCHEMA, **schema_updates}
-    return parse_meds_csvs(load_yaml(yaml_str, Loader=Loader), schema=schema)
+    return parse_meds_csvs(load_yaml(yaml_str.strip(), Loader=Loader), schema=schema)
 
 
 def dict_to_hydra_kwargs(d: dict[str, str]) -> str:
@@ -137,6 +142,9 @@ def run_command(
 
     err_cmd_lines = []
 
+    if config_name is not None and not config_name.startswith("_"):
+        config_name = f"_{config_name}"
+
     if do_use_config_yaml:
         if config_name is None:
             raise ValueError("config_name must be provided if do_use_config_yaml is True.")
@@ -271,6 +279,10 @@ def check_NRT_output(
 FILE_T = pl.DataFrame | dict[str, Any] | str
 
 
+def add_params(templ_str: str, **kwargs):
+    return templ_str.format(**kwargs)
+
+
 @contextmanager
 def input_dataset(input_files: dict[str, FILE_T] | None = None):
     with tempfile.TemporaryDirectory() as d:
@@ -294,6 +306,12 @@ def input_dataset(input_files: dict[str, FILE_T] | None = None):
                     fp.write_text(json.dumps(data))
                 case str():
                     fp.write_text(data.strip())
+                case _ if callable(data):
+                    data_str = data(
+                        input_dir=str(input_dir.resolve()),
+                        cohort_dir=str(cohort_dir.resolve()),
+                    )
+                    fp.write_text(data_str)
                 case _:
                     raise ValueError(f"Unknown data type {type(data)} for file {fp.relative_to(input_dir)}")
 
@@ -356,7 +374,7 @@ def check_outputs(
 
 def single_stage_tester(
     script: str | Path,
-    stage_name: str,
+    stage_name: str | None,
     stage_kwargs: dict[str, str] | None,
     do_pass_stage_name: bool = False,
     do_use_config_yaml: bool = False,
@@ -366,8 +384,15 @@ def single_stage_tester(
     config_name: str = "preprocess",
     input_files: dict[str, FILE_T] | None = None,
     df_check_kwargs: dict | None = None,
+    test_name: str | None = None,
+    do_include_dirs: bool = True,
+    hydra_verbose: bool = True,
+    stdout_regex: str | None = None,
     **pipeline_kwargs,
 ):
+    if test_name is None:
+        test_name = f"Single stage transform: {stage_name}"
+
     if df_check_kwargs is None:
         df_check_kwargs = {}
 
@@ -377,20 +402,23 @@ def single_stage_tester(
                 pipeline_kwargs[k] = v.format(input_dir=str(input_dir.resolve()))
 
         pipeline_config_kwargs = {
-            "input_dir": str(input_dir.resolve()),
-            "cohort_dir": str(cohort_dir.resolve()),
-            "stages": [stage_name],
-            "hydra.verbose": True,
+            "hydra.verbose": hydra_verbose,
             **pipeline_kwargs,
         }
 
+        if do_include_dirs:
+            pipeline_config_kwargs["input_dir"] = str(input_dir.resolve())
+            pipeline_config_kwargs["cohort_dir"] = str(cohort_dir.resolve())
+
+        if stage_name is not None:
+            pipeline_config_kwargs["stages"] = [stage_name]
         if stage_kwargs:
             pipeline_config_kwargs["stage_configs"] = {stage_name: stage_kwargs}
 
         run_command_kwargs = {
             "script": script,
             "hydra_kwargs": pipeline_config_kwargs,
-            "test_name": f"Single stage transform: {stage_name}",
+            "test_name": test_name,
             "should_error": should_error,
             "config_name": config_name,
             "do_use_config_yaml": do_use_config_yaml,
@@ -405,6 +433,12 @@ def single_stage_tester(
         if should_error:
             return
 
+        if stdout_regex is not None:
+            regex = re.compile(stdout_regex)
+            assert regex.search(stdout) is not None, (
+                f"Expected stdout to match regex:\n{stdout_regex}\n" f"Got:\n{stdout}"
+            )
+
         try:
             check_outputs(
                 cohort_dir,