split_and_shard_subjects

Extracts the set of unique subjects from the raw data and splits/shards them and saves the result.

This stage splits the subjects into training, tuning, and held-out sets, and further splits those sets into shards.

All arguments are specified through the command line into the cfg object through Hydra.

The cfg.stage_cfg object is a special key that is imputed by OmegaConf to contain the stage-specific configuration arguments based on the global, pipeline-level configuration file.

Details

Property	Value
Type	main
Metadata stage	True

Default Configuration

n_subjects_per_shard: 50000
external_splits_json_fp: null
split_fracs:
  train: 0.8
  tuning: 0.1
  held_out: 0.1

Usage

MEDS_transform-stage <pipeline.yaml> split_and_shard_subjects input_dir=<input> output_dir=<output>

Examples

join

Demonstrates split-and-shard over a joined table: the vitals table exposes subject_id only via a _table.join on stays. The split machinery follows that join to discover unique subject IDs before assigning splits.

Stage configuration:

split_fracs:
  train: 0.5
  tuning: 0.5
  held_out: null
n_subjects_per_shard: 10

This example uses the stage’s config.yaml file.

Input files:

data/vitals/[0-3).parquet:
  stay_id: [10, 10, 20]
  charttime:
    ["01/01/2021 00:00:00", "01/01/2021 01:00:00", "01/01/2021 02:00:00"]
  HR: [70, 75, 65]

data/stays/[0-2).parquet:
  stay_id: [10, 20]
  subject_id: [111, 222]

messy.yaml: |
  vitals:
    _table:
      join:
        stays:
          key: stay_id
          cols: [subject_id]
    HR:
      code: HR
      time: '$charttime::"%m/%d/%Y %H:%M:%S"'
      numeric_value: $HR

Expected output metadata:

metadata/.shards.json:
  train/0: [111]
  tuning/0: [222]

Run this stage:

MEDS_transform-stage <pipeline.yaml> split_and_shard_subjects input_dir=<input> output_dir=<output>

default

Scans unique subject_id values across the sub-sharded input tables and partitions them into train/tuning/held_out splits, then into shards of at most n_subjects_per_shard subjects per shard. The resulting split/shard map is written as metadata/.shards.json, which all downstream stages consume.

Since the dataset here has 4 subjects and the 50/25/25 split assigns subjects deterministically under seed: 1, each split gets at most one shard.

Stage configuration:

split_fracs:
  train: 0.5
  tuning: 0.25
  held_out: 0.25
n_subjects_per_shard: 10

This example uses the stage’s config.yaml file.

Input files:

data/patients/[0-2).parquet:
  MRN: [1, 2]
  dob: ["2000-01-01T00:00:00", "2001-02-02T00:00:00"]
  eye_color: ["BROWN", "BLUE"]

data/patients/[2-4).parquet:
  MRN: [3, 4]
  dob: ["2002-03-03T00:00:00", "2003-04-04T00:00:00"]
  eye_color: ["GREEN", "BROWN"]

data/labs/[0-2).parquet:
  patient_id: [1, 1]
  timestamp: ["2020-01-01T10:00:00", "2020-01-01T11:00:00"]
  test_name: ["HR", "TEMP"]
  result: [80.0, 36.6]

data/labs/[2-4).parquet:
  patient_id: [2, 2]
  timestamp: ["2020-01-02T12:00:00", "2020-01-02T13:00:00"]
  test_name: ["HR", "TEMP"]
  result: [75.0, 37.0]

data/labs/[4-6).parquet:
  patient_id: [3, 3]
  timestamp: ["2020-01-03T14:00:00", "2020-01-03T15:00:00"]
  test_name: ["HR", "TEMP"]
  result: [85.0, 36.5]

data/labs/[6-8).parquet:
  patient_id: [4, 4]
  timestamp: ["2020-01-04T09:00:00", "2020-01-04T10:00:00"]
  test_name: ["HR", "TEMP"]
  result: [70.0, 36.8]

messy.yaml: |
  _defaults:
    subject_id: $MRN
  patients:
    eye_color:
      code: 'f"EYE_COLOR//{$eye_color}"'
      time: null
    dob:
      code: MEDS_BIRTH
      time: '$dob::"%Y-%m-%dT%H:%M:%S"'
  labs:
    _defaults:
      subject_id: $patient_id
    lab:
      code: $test_name
      time: '$timestamp::"%Y-%m-%dT%H:%M:%S"'
      numeric_value: $result

Expected output metadata:

metadata/.shards.json:
  train/0: [1, 4]
  tuning/0: [3]
  held_out/0: [2]

Run this stage:

MEDS_transform-stage <pipeline.yaml> split_and_shard_subjects input_dir=<input> output_dir=<output>