merge_to_MEDS_cohort

Merges the subject sub-sharded events into a single parquet file per subject shard.

This function reads, for each shard, the per-table file <prefix>.parquet under the shard’s directory in cfg.stage_cfg.input_dir — one file per configured table prefix, in config order — and merges them into a single dataframe. All such dataframes are assumed to be in the unnested, MEDS format, and cover the same group of subjects (specific to the shard being processed). The merged dataframe will also be sorted by subject ID and time. The internal code_components struct column is dropped during the merge (#254; see :func:merge_subdirs_and_sort) — metadata extraction reads the pre-merge per-table events, so the merged data does not need it and unifying the per-table structs is catastrophically memory-expensive.

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	False

Default Configuration

unique_by: '*'
additional_sort_by: null

Usage

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

Examples

default

Concatenates the per-table event parquets (convert_to_MEDS_events output) within each (split, shard) into a single file at data/<split>/<shard>.parquet. Rows are sorted so null time values (static events like EYE_COLOR) precede real timestamps per subject. The internal code_components struct is not carried into the merged output.

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

Input files:

data/train/0/patients.parquet:
  subject_id: [1, 4, 1, 4]
  code: ["EYE_COLOR//BROWN", "EYE_COLOR//BROWN", "MEDS_BIRTH", "MEDS_BIRTH"]
  code_components:
    - { eye_color: BROWN }
    - { eye_color: BROWN }
    - null
    - null
  time:
    - null
    - null
    - 2000-01-01T00:00:00
    - 2003-04-04T00:00:00
  source_block:
    [patients/eye_color, patients/eye_color, patients/dob, patients/dob]

data/train/0/labs.parquet:
  subject_id: [1, 1, 4, 4]
  code: [HR, TEMP, HR, TEMP]
  code_components:
    - { test_name: HR }
    - { test_name: TEMP }
    - { test_name: HR }
    - { test_name: TEMP }
  time:
    [
      2020-01-01T10:00:00,
      2020-01-01T11:00:00,
      2020-01-04T09:00:00,
      2020-01-04T10:00:00,
    ]
  numeric_value: [80.0, 36.6, 70.0, 36.8]
  source_block: [labs/lab, labs/lab, labs/lab, labs/lab]

data/tuning/0/patients.parquet:
  subject_id: [3, 3]
  code: [EYE_COLOR//GREEN, MEDS_BIRTH]
  code_components:
    - { eye_color: GREEN }
    - null
  time:
    - null
    - 2002-03-03T00:00:00
  source_block: [patients/eye_color, patients/dob]

data/tuning/0/labs.parquet:
  subject_id: [3, 3]
  code: [HR, TEMP]
  code_components:
    - { test_name: HR }
    - { test_name: TEMP }
  time: [2020-01-03T14:00:00, 2020-01-03T15:00:00]
  numeric_value: [85.0, 36.5]
  source_block: [labs/lab, labs/lab]

data/held_out/0/patients.parquet:
  subject_id: [2, 2]
  code: [EYE_COLOR//BLUE, MEDS_BIRTH]
  code_components:
    - { eye_color: BLUE }
    - null
  time:
    - null
    - 2001-02-02T00:00:00
  source_block: [patients/eye_color, patients/dob]

data/held_out/0/labs.parquet:
  subject_id: [2, 2]
  code: [HR, TEMP]
  code_components:
    - { test_name: HR }
    - { test_name: TEMP }
  time: [2020-01-02T12:00:00, 2020-01-02T13:00:00]
  numeric_value: [75.0, 37.0]
  source_block: [labs/lab, labs/lab]

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

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 files:

data/train/0.parquet:
  subject_id: [1, 1, 1, 1, 4, 4, 4, 4]
  code:
    [
      EYE_COLOR//BROWN,
      MEDS_BIRTH,
      HR,
      TEMP,
      EYE_COLOR//BROWN,
      MEDS_BIRTH,
      HR,
      TEMP,
    ]
  time:
    - null
    - 2000-01-01T00:00:00
    - 2020-01-01T10:00:00
    - 2020-01-01T11:00:00
    - null
    - 2003-04-04T00:00:00
    - 2020-01-04T09:00:00
    - 2020-01-04T10:00:00
  source_block:
    - patients/eye_color
    - patients/dob
    - labs/lab
    - labs/lab
    - patients/eye_color
    - patients/dob
    - labs/lab
    - labs/lab
  numeric_value: [null, null, 80.0, 36.6, null, null, 70.0, 36.8]

data/tuning/0.parquet:
  subject_id: [3, 3, 3, 3]
  code: [EYE_COLOR//GREEN, MEDS_BIRTH, HR, TEMP]
  time:
    - null
    - 2002-03-03T00:00:00
    - 2020-01-03T14:00:00
    - 2020-01-03T15:00:00
  source_block: [patients/eye_color, patients/dob, labs/lab, labs/lab]
  numeric_value: [null, null, 85.0, 36.5]

data/held_out/0.parquet:
  subject_id: [2, 2, 2, 2]
  code: [EYE_COLOR//BLUE, MEDS_BIRTH, HR, TEMP]
  time:
    - null
    - 2001-02-02T00:00:00
    - 2020-01-02T12:00:00
    - 2020-01-02T13:00:00
  source_block: [patients/eye_color, patients/dob, labs/lab, labs/lab]
  numeric_value: [null, null, 75.0, 37.0]

Run this stage:

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