Repository: mmcdermott/MEDS_transforms

Issue:
Pipelines should automatically determine shards from the input directory rather than relying on the `splits.json` file.
This would allow pipelines to be used even on dataests not extracted with `MEDS-extract`. This entails:
  - [x] Cleaning up the interface to not rely on this and automatically load shards.
  - [x] Ensuring a clean, sensible solution for relying on either shard prefixes or data filters for `patient_splits.parquet` integration as needed.

Discussion:

Patch:
diff --git a/src/MEDS_transforms/configs/extract.yaml b/src/MEDS_transforms/configs/extract.yaml
index 377c39c..b0f2d50 100644
--- a/src/MEDS_transforms/configs/extract.yaml
+++ b/src/MEDS_transforms/configs/extract.yaml
@@ -30,7 +30,7 @@ event_conversion_config_fp: ???
 # The code modifier columns are in this pipeline only used in the aggregate_code_metadata stage.
 code_modifiers: null
 # The shards mapping is stored in the root of the final output directory.
-shards_map_fp: "${cohort_dir}/splits.json"
+shards_map_fp: "${cohort_dir}/metadata/.shards.json"
 
 stages:
   - shard_events
diff --git a/src/MEDS_transforms/configs/preprocess.yaml b/src/MEDS_transforms/configs/preprocess.yaml
index 8a30b43..6b2dc8b 100644
--- a/src/MEDS_transforms/configs/preprocess.yaml
+++ b/src/MEDS_transforms/configs/preprocess.yaml
@@ -21,9 +21,6 @@ etl_metadata.pipeline_name: "preprocess"
 #    tokenization.
 code_modifiers: ???
 
-# The shards map filepath is stored in the global input directory for model-specific pre-processing.
-shards_map_fp: "${input_dir}/splits.json"
-
 # Pipeline Structure
 stages:
   - filter_patients
diff --git a/src/MEDS_transforms/configs/stage_configs/split_and_shard_patients.yaml b/src/MEDS_transforms/configs/stage_configs/split_and_shard_patients.yaml
index b56aeb1..c4015bd 100644
--- a/src/MEDS_transforms/configs/stage_configs/split_and_shard_patients.yaml
+++ b/src/MEDS_transforms/configs/stage_configs/split_and_shard_patients.yaml
@@ -1,6 +1,6 @@
 split_and_shard_patients:
   is_metadata: True
-  output_dir: ${cohort_dir}
+  output_dir: ${cohort_dir}/metadata
   n_patients_per_shard: 50000
   external_splits_json_fp: null
   split_fracs:
diff --git a/src/MEDS_transforms/extract/convert_to_sharded_events.py b/src/MEDS_transforms/extract/convert_to_sharded_events.py
index e5eb514..8d6a01e 100755
--- a/src/MEDS_transforms/extract/convert_to_sharded_events.py
+++ b/src/MEDS_transforms/extract/convert_to_sharded_events.py
@@ -685,9 +685,9 @@ def main(cfg: DictConfig):
     file.
     """
 
-    input_dir, patient_subsharded_dir, metadata_input_dir, shards_map_fn = stage_init(cfg)
+    input_dir, patient_subsharded_dir, metadata_input_dir = stage_init(cfg)
 
-    shards = json.loads(shards_map_fn.read_text())
+    shards = json.loads(Path(cfg.shards_map_fp).read_text())
 
     event_conversion_cfg_fp = Path(cfg.event_conversion_config_fp)
     if not event_conversion_cfg_fp.exists():
diff --git a/src/MEDS_transforms/extract/extract_code_metadata.py b/src/MEDS_transforms/extract/extract_code_metadata.py
index 5e831be..db08f7b 100644
--- a/src/MEDS_transforms/extract/extract_code_metadata.py
+++ b/src/MEDS_transforms/extract/extract_code_metadata.py
@@ -349,7 +349,7 @@ def main(cfg: DictConfig):
             schema.
     """
 
-    stage_input_dir, partial_metadata_dir, _, _ = stage_init(cfg)
+    stage_input_dir, partial_metadata_dir, _ = stage_init(cfg)
     raw_input_dir = Path(cfg.input_dir)
 
     event_conversion_cfg_fp = Path(cfg.event_conversion_config_fp)
diff --git a/src/MEDS_transforms/extract/finalize_MEDS_metadata.py b/src/MEDS_transforms/extract/finalize_MEDS_metadata.py
index f8486c2..366d89a 100755
--- a/src/MEDS_transforms/extract/finalize_MEDS_metadata.py
+++ b/src/MEDS_transforms/extract/finalize_MEDS_metadata.py
@@ -146,7 +146,7 @@ def main(cfg: DictConfig):
         logger.info("Non-zero worker found in reduce-only stage. Exiting")
         return
 
-    _, _, input_metadata_dir, shards_map_fp = stage_init(cfg)
+    _, _, input_metadata_dir = stage_init(cfg)
     output_metadata_dir = Path(cfg.stage_cfg.reducer_output_dir)
 
     output_code_metadata_fp = output_metadata_dir / "codes.parquet"
@@ -193,6 +193,7 @@ def main(cfg: DictConfig):
     dataset_metadata_fp.write_text(json.dumps(dataset_metadata))
 
     # Split creation
+    shards_map_fp = Path(cfg.shards_map_fp)
     logger.info("Creating patient splits from {str(shards_map_fp.resolve())}")
     shards_map = json.loads(shards_map_fp.read_text())
     patient_splits = []
diff --git a/src/MEDS_transforms/extract/merge_to_MEDS_cohort.py b/src/MEDS_transforms/extract/merge_to_MEDS_cohort.py
index b9e638a..df8f301 100755
--- a/src/MEDS_transforms/extract/merge_to_MEDS_cohort.py
+++ b/src/MEDS_transforms/extract/merge_to_MEDS_cohort.py
@@ -1,4 +1,6 @@
 #!/usr/bin/env python
+import json
+import random
 from functools import partial
 from pathlib import Path
 
@@ -8,7 +10,7 @@ from loguru import logger
 from omegaconf import DictConfig, OmegaConf
 
 from MEDS_transforms.extract import CONFIG_YAML
-from MEDS_transforms.mapreduce.mapper import map_over, shard_iterator
+from MEDS_transforms.mapreduce.mapper import map_over
 
 
 def merge_subdirs_and_sort(
@@ -234,10 +236,37 @@ def main(cfg: DictConfig):
         additional_sort_by=cfg.stage_cfg.get("additional_sort_by", None),
     )
 
+    shard_map_fp = Path(cfg.shards_map_fp)
+    if not shard_map_fp.exists():
+        raise FileNotFoundError(f"Shard map file not found at {str(shard_map_fp.resolve())}")
+
+    shards = list(json.loads(shard_map_fp.read_text()).keys())
+
+    def shard_iterator(cfg: DictConfig) -> tuple[list[str], bool]:
+        input_dir = Path(cfg.stage_cfg.data_input_dir)
+        output_dir = Path(cfg.stage_cfg.output_dir)
+
+        if cfg.stage_cfg.get("train_only", None):
+            raise ValueError("train_only is not supported for this stage.")
+
+        if "worker" in cfg:
+            random.seed(cfg.worker)
+        random.shuffle(shards)
+
+        logger.info(f"Mapping computation over a maximum of {len(shards)} shards")
+
+        out = []
+        for sh in shards:
+            in_fp = input_dir / sh
+            out_fp = output_dir / f"{sh}.parquet"
+            out.append((in_fp, out_fp))
+
+        return out, False
+
     map_over(
         cfg,
         read_fn=read_fn,
-        shard_iterator_fntr=partial(shard_iterator, in_suffix=""),
+        shard_iterator_fntr=shard_iterator,
     )
 
 
diff --git a/src/MEDS_transforms/extract/split_and_shard_patients.py b/src/MEDS_transforms/extract/split_and_shard_patients.py
index 3a4e1d5..ab45a64 100755
--- a/src/MEDS_transforms/extract/split_and_shard_patients.py
+++ b/src/MEDS_transforms/extract/split_and_shard_patients.py
@@ -186,7 +186,7 @@ def main(cfg: DictConfig):
             ensure that split fractions sum to 1.
     """
 
-    subsharded_dir, MEDS_cohort_dir, _, _ = stage_init(cfg)
+    subsharded_dir, _, _ = stage_init(cfg)
 
     event_conversion_cfg_fp = Path(cfg.event_conversion_config_fp)
     if not event_conversion_cfg_fp.exists():
@@ -249,11 +249,11 @@ def main(cfg: DictConfig):
         seed=cfg.seed,
     )
 
-    logger.info(f"Writing sharded patients to {MEDS_cohort_dir}")
-    MEDS_cohort_dir.mkdir(parents=True, exist_ok=True)
-    out_fp = MEDS_cohort_dir / "splits.json"
-    out_fp.write_text(json.dumps(sharded_patients))
-    logger.info(f"Done writing sharded patients to {out_fp}")
+    shards_map_fp = Path(cfg.shards_map_fp)
+    logger.info(f"Writing sharded patients to {str(shards_map_fp.resolve())}")
+    shards_map_fp.parent.mkdir(parents=True, exist_ok=True)
+    shards_map_fp.write_text(json.dumps(sharded_patients))
+    logger.info("Done writing sharded patients")
 
 
 if __name__ == "__main__":
diff --git a/src/MEDS_transforms/mapreduce/mapper.py b/src/MEDS_transforms/mapreduce/mapper.py
index 3d675b0..6b637f3 100644
--- a/src/MEDS_transforms/mapreduce/mapper.py
+++ b/src/MEDS_transforms/mapreduce/mapper.py
@@ -157,27 +157,33 @@ def map_over(
 
     start = datetime.now()
 
-    process_split = cfg.stage_cfg.get("process_split", None)
+    train_only = cfg.stage_cfg.get("train_only", False)
     split_fp = Path(cfg.input_dir) / "metadata" / "patient_split.parquet"
-    shards_map_fp = Path(cfg.shards_map_fp) if "shards_map_fp" in cfg else None
-    if process_split and split_fp.exists():
-        split_patients = (
-            pl.scan_parquet(split_fp)
-            .filter(pl.col("split") == process_split)
-            .select(pl.col("patient_id"))
-            .collect()
-            .to_list()
-        )
-        read_fn = read_and_filter_fntr(split_patients, read_fn)
-    elif process_split and shards_map_fp and shards_map_fp.exists():
-        logger.warning(
-            f"Split {process_split} requested, but no patient split file found at {str(split_fp)}. "
-            f"Assuming this is handled through shard filtering."
-        )
-    elif process_split:
-        raise ValueError(
-            f"Split {process_split} requested, but no patient split file found at {str(split_fp)}."
-        )
+
+    shards, includes_only_train = shard_iterator_fntr(cfg)
+
+    if train_only:
+        if includes_only_train:
+            logger.info(
+                f"Processing train split only via shard prefix. Not filtering with {str(split_fp.resolve())}."
+            )
+        elif split_fp.exists():
+            logger.info(f"Processing train split only by filtering read dfs via {str(split_fp.resolve())}")
+            train_patients = (
+                pl.scan_parquet(split_fp)
+                .filter(pl.col("split") == "train")
+                .select(pl.col("patient_id"))
+                .collect()
+                .to_list()
+            )
+            read_fn = read_and_filter_fntr(train_patients, read_fn)
+        else:
+            raise FileNotFoundError(
+                f"Train split requested, but shard prefixes can't be used and "
+                f"patient split file not found at {str(split_fp.resolve())}."
+            )
+    elif includes_only_train:
+        raise ValueError("All splits should be used, but shard iterator is returning only train splits?!?")
 
     def fntr_params(compute_fn: ANY_COMPUTE_FN_T) -> ComputeFnArgs:
         compute_fn_params = inspect.signature(compute_fn).parameters
@@ -209,7 +215,7 @@ def map_over(
             raise ValueError("Invalid compute function")
 
     all_out_fps = []
-    for in_fp, out_fp in shard_iterator_fntr(cfg):
+    for in_fp, out_fp in shards:
         logger.info(f"Processing {str(in_fp.resolve())} into {str(out_fp.resolve())}")
         rwlock_wrap(
             in_fp,
diff --git a/src/MEDS_transforms/mapreduce/utils.py b/src/MEDS_transforms/mapreduce/utils.py
index 77cbe7d..e94df8d 100644
--- a/src/MEDS_transforms/mapreduce/utils.py
+++ b/src/MEDS_transforms/mapreduce/utils.py
@@ -1,7 +1,7 @@
 """Basic utilities for parallelizable mapreduces on sharded MEDS datasets with caching and locking."""
 
+import hashlib
 import json
-import random
 from collections.abc import Callable
 from datetime import datetime
 from pathlib import Path
@@ -223,122 +223,226 @@ def rwlock_wrap[
 
 def shard_iterator(
     cfg: DictConfig,
-    in_suffix: str = ".parquet",
     out_suffix: str = ".parquet",
     in_prefix: str = "",
-    out_prefix: str = "",
-):
-    """Provides a generator that yields shard input and output files for mapreduce operations.
+) -> tuple[list[tuple[Path, Path]], bool]:
+    """Returns a list of the shards found in the input directory and their corresponding output directories.
 
     Args:
         cfg: The configuration dictionary for the overall pipeline. Should (possibly) contain the following
             keys (some are optional, as marked below):
-            - ``stage_cfg.data_input_dir`` (mandatory): The directory containing the input data.
-            - ``stage_cfg.output_dir`` (mandatory): The directory to write the output data.
-            - ``shards_map_fp`` (mandatory): The file path to the shards map JSON file.
-            - ``stage_cfg.process_split`` (optional): The prefix of the shards to process (e.g.,
-              ``"train/"``). If not provided, all shards will be processed.
-            - ``worker`` (optional): The worker ID for the MR worker; this is also used to seed the
-              randomization process. If not provided, the randomization process is unseeded.
-        in_suffix: The suffix of the input files. Defaults to ".parquet". This can be set to "" to process
-            entire directories.
+            - `stage_cfg.data_input_dir` (mandatory): The directory containing the input data.
+            - `stage_cfg.output_dir` (mandatory): The directory to write the output data.
+            - `stage_cfg.train_only` (optional): The prefix of the shards to process (e.g.,
+              `"train/"`). If not provided, all shards will be processed.
+            - `worker` (optional): The worker ID for the MR worker; this is also used to seed the
+              randomization process. If not provided, the randomization process will be unseeded.
         out_suffix: The suffix of the output files. Defaults to ".parquet".
-        in_prefix: The prefix of the input files. Defaults to "". This can be used to load files from a
-            subdirectory of the input directory by including a "/" at the end of the prefix.
-        out_prefix: The prefix of the output files. Defaults to "".
+        in_prefix: The prefix of the input files. Defaults to "". This must be a full path component. It can
+            end with a slash but even if it doesn't it will be interpreted as a full path component.
 
     Yields:
         Randomly shuffled pairs of input and output file paths for each shard. The randomization process is
         seeded by the worker ID in ``cfg``, if provided, otherwise it is left unseeded.
 
     Examples:
-        >>> from tempfile import NamedTemporaryFile
-        >>> shards = {"train/0": [1, 2, 3], "train/1": [4, 5, 6], "held_out": [4, 5, 6], "foo": [5]}
-        >>> with NamedTemporaryFile() as tmp:
-        ...     _ = Path(tmp.name).write_text(json.dumps(shards))
+        >>> from tempfile import TemporaryDirectory
+        >>> import polars as pl
+        >>> df = pl.DataFrame({
+        ...     "patient_id": [1, 2, 3, 4, 5, 6, 7, 8, 9],
+        ...     "code": ["A", "B", "C", "D", "E", "F", "G", "H", "I"],
+        ...     "time": [1, 2, 3, 4, 5, 6, 1, 2, 3],
+        ... })
+        >>> shards = {"train/0": [1, 2, 3, 4], "train/1": [5, 6, 7], "tuning": [8], "held_out": [9]}
+        >>> def write_dfs(input_dir: Path, df: pl.DataFrame=df, shards: dict=shards, sfx: str=".parquet"):
+        ...     for shard_name, patient_ids in shards.items():
+        ...         df = df.filter(pl.col("patient_id").is_in(patient_ids))
+        ...         shard_fp = input_dir / f"{shard_name}{sfx}"
+        ...         shard_fp.parent.mkdir(exist_ok=True, parents=True)
+        ...         if sfx == ".parquet": df.write_parquet(shard_fp)
+        ...         elif sfx == ".csv": df.write_csv(shard_fp)
+        ...         else: raise ValueError(f"Unsupported suffix {sfx}")
+        ...     return
+
+    By default, this will load all shards in the input directory and write specify their appropriate output
+    directories:
+        >>> with TemporaryDirectory() as tmp:
+        ...     root = Path(tmp)
+        ...     input_dir = root / "data"
+        ...     output_dir = root / "output"
+        ...     write_dfs(input_dir)
+        ...     cfg = DictConfig({
+        ...         "stage_cfg": {"data_input_dir": str(input_dir), "output_dir": str(output_dir)},
+        ...         "worker": 1,
+        ...     })
+        ...     fps, includes_only_train = shard_iterator(cfg)
+        >>> [(i.relative_to(root), o.relative_to(root)) for i, o in fps] # doctest: +NORMALIZE_WHITESPACE
+        [(PosixPath('data/train/1.parquet'),  PosixPath('output/train/1.parquet')),
+         (PosixPath('data/held_out.parquet'), PosixPath('output/held_out.parquet')),
+         (PosixPath('data/tuning.parquet'),   PosixPath('output/tuning.parquet')),
+         (PosixPath('data/train/0.parquet'),  PosixPath('output/train/0.parquet'))]
+        >>> includes_only_train
+        False
+
+    Different workers will shuffle the shards differently:
+        >>> with TemporaryDirectory() as tmp:
+        ...     root = Path(tmp)
+        ...     input_dir = root / "data"
+        ...     output_dir = root / "output"
+        ...     write_dfs(input_dir)
         ...     cfg = DictConfig({
-        ...         "stage_cfg": {"data_input_dir": "data/", "output_dir": "output/"},
-        ...         "shards_map_fp": tmp.name, "worker": 1,
+        ...         "stage_cfg": {"data_input_dir": str(input_dir), "output_dir": str(output_dir)},
+        ...         "worker": 2,
         ...     })
-        ...     gen = shard_iterator(cfg)
-        ...     list(gen) # doctest: +NORMALIZE_WHITESPACE
-        [(PosixPath('data/foo.parquet'),      PosixPath('output/foo.parquet')),
-         (PosixPath('data/train/0.parquet'),  PosixPath('output/train/0.parquet')),
+        ...     fps, includes_only_train = shard_iterator(cfg)
+        >>> [(i.relative_to(root), o.relative_to(root)) for i, o in fps] # doctest: +NORMALIZE_WHITESPACE
+        [(PosixPath('data/tuning.parquet'),   PosixPath('output/tuning.parquet')),
          (PosixPath('data/held_out.parquet'), PosixPath('output/held_out.parquet')),
-         (PosixPath('data/train/1.parquet'),  PosixPath('output/train/1.parquet'))]
-        >>> with NamedTemporaryFile() as tmp:
-        ...     _ = Path(tmp.name).write_text(json.dumps(shards))
+         (PosixPath('data/train/1.parquet'),  PosixPath('output/train/1.parquet')),
+         (PosixPath('data/train/0.parquet'),  PosixPath('output/train/0.parquet'))]
+        >>> includes_only_train
+        False
+
+    We can also make it look within a specific input subdir of the data directory and change the output
+    suffix. Note that using a specific input subdir is _different_ than requesting it load only train.
+        >>> with TemporaryDirectory() as tmp:
+        ...     root = Path(tmp)
+        ...     input_dir = root / "data"
+        ...     output_dir = root / "output"
+        ...     write_dfs(input_dir)
         ...     cfg = DictConfig({
-        ...         "stage_cfg": {"data_input_dir": "data/", "output_dir": "output/"},
-        ...         "shards_map_fp": tmp.name, "worker": 1,
+        ...         "stage_cfg": {"data_input_dir": str(input_dir), "output_dir": str(output_dir)},
+        ...         "worker": 1,
         ...     })
-        ...     gen = shard_iterator(cfg, in_suffix="", out_suffix=".csv", in_prefix="a/", out_prefix="b/")
-        ...     list(gen) # doctest: +NORMALIZE_WHITESPACE
-        [(PosixPath('data/a/foo'),      PosixPath('output/b/foo.csv')),
-         (PosixPath('data/a/train/0'),  PosixPath('output/b/train/0.csv')),
-         (PosixPath('data/a/held_out'), PosixPath('output/b/held_out.csv')),
-         (PosixPath('data/a/train/1'),  PosixPath('output/b/train/1.csv'))]
-        >>> with NamedTemporaryFile() as tmp:
-        ...     _ = Path(tmp.name).write_text(json.dumps(shards))
+        ...     fps, includes_only_train = shard_iterator(cfg, in_prefix="train", out_suffix=".csv")
+        >>> [(i.relative_to(root), o.relative_to(root)) for i, o in fps] # doctest: +NORMALIZE_WHITESPACE
+        [(PosixPath('data/train/0.parquet'),  PosixPath('output/0.csv')),
+         (PosixPath('data/train/1.parquet'),  PosixPath('output/1.csv'))]
+        >>> includes_only_train
+        False
+
+    We can also make it load only 'train' shards, in the case that there are shards with a valid "train/"
+    prefix.
+        >>> with TemporaryDirectory() as tmp:
+        ...     root = Path(tmp)
+        ...     input_dir = root / "data"
+        ...     output_dir = root / "output"
+        ...     write_dfs(input_dir)
         ...     cfg = DictConfig({
         ...         "stage_cfg": {
-        ...             "data_input_dir": "data/", "output_dir": "output/", "process_split": "train/"
+        ...             "data_input_dir": str(input_dir), "output_dir": str(output_dir),
+        ...             "train_only": True,
         ...         },
-        ...         "shards_map_fp": tmp.name,
         ...         "worker": 1,
         ...     })
-        ...     gen = shard_iterator(cfg)
-        ...     list(gen) # doctest: +NORMALIZE_WHITESPACE
+        ...     fps, includes_only_train = shard_iterator(cfg)
+        >>> [(i.relative_to(root), o.relative_to(root)) for i, o in fps] # doctest: +NORMALIZE_WHITESPACE
         [(PosixPath('data/train/1.parquet'),  PosixPath('output/train/1.parquet')),
          (PosixPath('data/train/0.parquet'),  PosixPath('output/train/0.parquet'))]
+        >>> includes_only_train
+        True
 
-    Note that if you ask it to process a split that isn't a valid prefix of the shards, it will return all
-    shards and assume that is covered via a `patient_splits.parquet` file:
-        >>> with NamedTemporaryFile() as tmp:
-        ...     _ = Path(tmp.name).write_text(json.dumps(shards))
+    The train prefix used is precisely `train/` -- other uses of train will not work:
+        >>> wrong_pfx_shards = {"train": [1, 2, 3], "train_1": [4, 5, 6], "train-2": [7, 8, 9]}
+        >>> with TemporaryDirectory() as tmp:
+        ...     root = Path(tmp)
+        ...     input_dir = root / "data"
+        ...     output_dir = root / "output"
+        ...     write_dfs(input_dir, shards=wrong_pfx_shards)
         ...     cfg = DictConfig({
-        ...         "stage_cfg": {"data_input_dir": "data/", "output_dir": "output/"},
-        ...         "shards_map_fp": tmp.name, "process_split": "nonexisting", "worker": 1,
+        ...         "stage_cfg": {
+        ...             "data_input_dir": str(input_dir), "output_dir": str(output_dir),
+        ...             "train_only": True,
+        ...         },
+        ...         "worker": 1,
         ...     })
-        ...     gen = shard_iterator(cfg)
-        ...     list(gen) # doctest: +NORMALIZE_WHITESPACE
-        [(PosixPath('data/foo.parquet'),      PosixPath('output/foo.parquet')),
-         (PosixPath('data/train/0.parquet'),  PosixPath('output/train/0.parquet')),
-         (PosixPath('data/held_out.parquet'), PosixPath('output/held_out.parquet')),
-         (PosixPath('data/train/1.parquet'),  PosixPath('output/train/1.parquet'))]
+        ...     fps, includes_only_train = shard_iterator(cfg)
+        >>> [(i.relative_to(root), o.relative_to(root)) for i, o in fps] # doctest: +NORMALIZE_WHITESPACE
+        [(PosixPath('data/train_1.parquet'),  PosixPath('output/train_1.parquet')),
+         (PosixPath('data/train-2.parquet'),  PosixPath('output/train-2.parquet')),
+         (PosixPath('data/train.parquet'),  PosixPath('output/train.parquet'))]
+        >>> includes_only_train
+        False
+
+    If there are no such shards, then it loads them all and assumes the filtering will be handled via the
+    splits parquet file.
+        >>> no_pfx_shards = {"0": [1, 2, 3], "1": [4, 5, 6], "2": [7, 8, 9]}
+        >>> with TemporaryDirectory() as tmp:
+        ...     root = Path(tmp)
+        ...     input_dir = root / "data"
+        ...     output_dir = root / "output"
+        ...     write_dfs(input_dir, shards=no_pfx_shards)
+        ...     cfg = DictConfig({
+        ...         "stage_cfg": {
+        ...             "data_input_dir": str(input_dir), "output_dir": str(output_dir),
+        ...             "train_only": True,
+        ...         },
+        ...         "worker": 1,
+        ...     })
+        ...     fps, includes_only_train = shard_iterator(cfg)
+        >>> [(i.relative_to(root), o.relative_to(root)) for i, o in fps] # doctest: +NORMALIZE_WHITESPACE
+        [(PosixPath('data/0.parquet'), PosixPath('output/0.parquet')),
+         (PosixPath('data/1.parquet'), PosixPath('output/1.parquet')),
+         (PosixPath('data/2.parquet'), PosixPath('output/2.parquet'))]
+        >>> includes_only_train
+        False
+
+    If it can't find any files, it will return an empty list:
+        >>> fps, includes_only_train = shard_iterator(cfg)
+        >>> fps
+        []
     """
 
     input_dir = Path(cfg.stage_cfg.data_input_dir)
     output_dir = Path(cfg.stage_cfg.output_dir)
-    shards_map_fp = Path(cfg.shards_map_fp) if "shards_map_fp" in cfg else None
 
-    if shards_map_fp and shards_map_fp.is_file():
-        shards = json.loads(shards_map_fp.read_text())
+    in_suffix = ".parquet"
 
-        if "process_split" in cfg.stage_cfg:
-            if any(k.startswith(cfg.stage_cfg.process_split) for k in shards):
-                logger.info(f'Processing shards with prefix "{cfg.stage_cfg.process_split}"')
-                shards = {k: v for k, v in shards.items() if k.startswith(cfg.stage_cfg.process_split)}
-        shards = list(shards.keys())
-    else:
-        shards = []
-        for p in input_dir.glob(f"{in_prefix}*{in_suffix}"):
-            relative_path = p.relative_to(input_dir)
-            shard_name = str(relative_path)
-            shard_name = shard_name[len(in_prefix) :] if in_prefix else shard_name
-            shard_name = shard_name[: -len(in_suffix)] if in_suffix else shard_name
-            shards.append(shard_name)
+    if in_prefix:
+        input_dir = input_dir / in_prefix
+
+    shards = []
+    for p in input_dir.glob(f"**/*{in_suffix}"):
+        relative_path = p.relative_to(input_dir)
+        shard_name = str(relative_path)
+        shard_name = shard_name[: -len(in_suffix)]
+        shards.append(shard_name)
+
+    # We initialize this to False and overwrite it if we find dedicated train shards.
+    includes_only_train = False
+
+    train_only = cfg.stage_cfg.get("train_only", None)
+    train_shards = [shard_name for shard_name in shards if shard_name.startswith("train/")]
+    if train_only and train_shards:
+        shards = train_shards
+        includes_only_train = True
+    elif train_only:
+        logger.info(
+            f"train_only={train_only} requested but no dedicated train shards found; processing all shards "
+            "and relying on `patient_splits.parquet` for filtering."
+        )
 
     if "worker" in cfg:
-        random.seed(cfg.worker)
-    random.shuffle(shards)
+        add_str = str(cfg.worker)
+    else:
+        add_str = str(datetime.now())
+
+    shard_keys = []
+    for shard in shards:
+        shard_hash = hashlib.sha256((add_str + shard).encode("utf-8")).hexdigest()
+        if shard_hash in shard_keys:
+            raise ValueError(f"Hash collision for shard {shard} with add_str {add_str}!")
+        shard_keys.append(int(shard_hash, 16))
+
+    shards = [shard for _, shard in sorted(zip(shard_keys, shards))]
 
     logger.info(f"Mapping computation over a maximum of {len(shards)} shards")
 
+    out = []
     for sp in shards:
-        in_fp = input_dir / f"{in_prefix}{sp}{in_suffix}"
-        out_fp = output_dir / f"{out_prefix}{sp}{out_suffix}"
-
+        in_fp = input_dir / f"{sp}{in_suffix}"
+        out_fp = output_dir / f"{sp}{out_suffix}"
         # TODO: Could add checking logic for existence of in_fp and/or out_fp here.
+        out.append((in_fp, out_fp))
 
-        yield in_fp, out_fp
+    return out, includes_only_train
diff --git a/src/MEDS_transforms/utils.py b/src/MEDS_transforms/utils.py
index bc7f036..0c4c20a 100644
--- a/src/MEDS_transforms/utils.py
+++ b/src/MEDS_transforms/utils.py
@@ -9,7 +9,6 @@ from typing import Any
 import hydra
 import polars as pl
 from loguru import logger
-from meds import train_split
 from omegaconf import DictConfig, OmegaConf
 
 from MEDS_transforms import __package_name__ as package_name
@@ -59,15 +58,14 @@ def write_lazyframe(df: pl.LazyFrame, out_fp: Path) -> None:
     df.write_parquet(out_fp, use_pyarrow=True)
 
 
-def stage_init(cfg: DictConfig):
+def stage_init(cfg: DictConfig) -> tuple[Path, Path, Path]:
     """Initializes the stage by logging the configuration and the stage-specific paths.
 
     Args:
         cfg: The global configuration object, which should have a ``cfg.stage_cfg`` attribute containing the
             stage specific configuration.
 
-    Returns: The data input directory, stage output directory, metadata input directory, and the shards file
-        path.
+    Returns: The data input directory, stage output directory, and metadata input directory.
     """
     hydra_loguru_init()
 
@@ -78,7 +76,6 @@ def stage_init(cfg: DictConfig):
     input_dir = Path(cfg.stage_cfg.data_input_dir)
     output_dir = Path(cfg.stage_cfg.output_dir)
     metadata_input_dir = Path(cfg.stage_cfg.metadata_input_dir)
-    shards_map_fp = Path(cfg.shards_map_fp)
 
     def chk(x: Path):
         return "✅" if x.exists() else "❌"
@@ -89,7 +86,6 @@ def stage_init(cfg: DictConfig):
             "input_dir": input_dir,
             "output_dir": output_dir,
             "metadata_input_dir": metadata_input_dir,
-            "shards_map_fp": shards_map_fp,
         }.items()
     ]
 
@@ -99,7 +95,7 @@ def stage_init(cfg: DictConfig):
     ]
     logger.debug("\n".join(logger_strs + paths_strs))
 
-    return input_dir, output_dir, metadata_input_dir, shards_map_fp
+    return input_dir, output_dir, metadata_input_dir
 
 
 def get_package_name() -> str:
@@ -226,7 +222,7 @@ def populate_stage(
         ...         "stage2": {"is_metadata": True},
         ...         "stage3": {"is_metadata": None, "output_dir": "/g/h"},
         ...         "stage4": {"data_input_dir": "/e/f"},
-        ...         "stage5": {"aggregations": ["foo"], "process_split": None},
+        ...         "stage5": {"aggregations": ["foo"], "train_only": None},
         ...     },
         ... })
         >>> args = [root_config[k] for k in ["input_dir", "cohort_dir", "stages", "stage_configs"]]
@@ -235,7 +231,7 @@ def populate_stage(
          'output_dir': '/c/d/stage1', 'reducer_output_dir': None}
         >>> populate_stage("stage2", *args) # doctest: +NORMALIZE_WHITESPACE
         {'is_metadata': True, 'data_input_dir': '/c/d/stage1', 'metadata_input_dir': '/a/b/metadata',
-         'output_dir': '/c/d/stage2', 'reducer_output_dir': '/c/d/stage2', 'process_split': 'train'}
+         'output_dir': '/c/d/stage2', 'reducer_output_dir': '/c/d/stage2', 'train_only': True}
         >>> populate_stage("stage3", *args) # doctest: +NORMALIZE_WHITESPACE
         {'is_metadata': None, 'output_dir': '/g/h', 'data_input_dir': '/c/d/stage1',
          'metadata_input_dir': '/c/d/stage2', 'reducer_output_dir': None}
@@ -243,7 +239,7 @@ def populate_stage(
         {'data_input_dir': '/e/f', 'is_metadata': False,
          'metadata_input_dir': '/c/d/stage2', 'output_dir': '/c/d/stage4', 'reducer_output_dir': None}
         >>> populate_stage("stage5", *args) # doctest: +NORMALIZE_WHITESPACE
-        {'aggregations': ['foo'], 'process_split': None, 'is_metadata': True, 'data_input_dir': '/c/d/stage4',
+        {'aggregations': ['foo'], 'train_only': None, 'is_metadata': True, 'data_input_dir': '/c/d/stage4',
          'metadata_input_dir': '/c/d/stage2', 'output_dir': '/c/d/stage5',
          'reducer_output_dir': '/c/d/metadata'}
         >>> populate_stage("stage6", *args) # doctest: +NORMALIZE_WHITESPACE
@@ -353,7 +349,7 @@ def populate_stage(
     }
 
     if is_metadata:
-        inferred_keys["process_split"] = train_split
+        inferred_keys["train_only"] = True
 
     out = {**stage}
     for key, val in inferred_keys.items():


Test patch:
diff --git a/tests/test_extract.py b/tests/test_extract.py
index f63b9c0..b59ae02 100644
--- a/tests/test_extract.py
+++ b/tests/test_extract.py
@@ -441,10 +441,10 @@ def test_extraction():
         all_stdouts.append(stdout)
 
         try:
-            splits_fp = MEDS_cohort_dir / "splits.json"
-            assert splits_fp.is_file(), f"Expected splits @ {str(splits_fp.resolve())} to exist."
+            shards_fp = MEDS_cohort_dir / "metadata" / ".shards.json"
+            assert shards_fp.is_file(), f"Expected splits @ {str(shards_fp.resolve())} to exist."
 
-            splits = json.loads(splits_fp.read_text())
+            splits = json.loads(shards_fp.read_text())
             expected_keys = ["train/0", "train/1", "tuning/0", "held_out/0"]
 
             expected_keys_str = ", ".join(f"'{k}'" for k in expected_keys)