Repository: mmcdermott/MEDS_transforms

Issue:
`parent_code` needs to be a list of strings in the code metadata schema and description a string.

Discussion:

Patch:
diff --git a/MIMIC-IV_Example/configs/event_configs.yaml b/MIMIC-IV_Example/configs/event_configs.yaml
index 7dc803a..f625d97 100644
--- a/MIMIC-IV_Example/configs/event_configs.yaml
+++ b/MIMIC-IV_Example/configs/event_configs.yaml
@@ -48,7 +48,7 @@ hosp/diagnoses_icd:
     _metadata:
       hosp/d_icd_diagnoses:
         description: "long_title"
-        parent_code: "ICD{icd_version}CM/{icd_code}" # Single strings are templates of columns.
+        parent_codes: "ICD{icd_version}CM/{icd_code}" # Single strings are templates of columns.
 
 hosp/drgcodes:
   drg:
@@ -105,7 +105,7 @@ hosp/labevents:
       d_labitems_to_loinc:
         description: ["omop_concept_name", "label"] # List of strings are columns to be collated
         itemid: "itemid (omop_source_code)"
-        parent_code: "{omop_vocabulary_id}/{omop_concept_code}"
+        parent_codes: "{omop_vocabulary_id}/{omop_concept_code}"
 
 hosp/omr:
   omr:
@@ -169,7 +169,7 @@ hosp/procedures_icd:
     _metadata:
       hosp/d_icd_procedures:
         description: "long_title"
-        parent_code: # List of objects are string labels mapping to filters to be evaluated.
+        parent_codes: # List of objects are string labels mapping to filters to be evaluated.
           - "ICD{icd_version}Proc/{icd_code}": { icd_version: 9 }
           - "ICD{icd_version}PCS/{icd_code}": { icd_version: 10 }
 
@@ -218,12 +218,12 @@ icu/chartevents:
       #meas_chartevents_main:
       #  description: ["omop_concept_name", "label"] # List of strings are columns to be collated
       #  itemid: "itemid (omop_source_code)"
-      #  parent_code: "{omop_vocabulary_id}/{omop_concept_code}"
+      #  parent_codes: "{omop_vocabulary_id}/{omop_concept_code}"
       # TODO: I don't know if this is necessary...
       d_labitems_to_loinc:
         description: ["omop_concept_name", "label"] # List of strings are columns to be collated
         itemid: "itemid (omop_source_code)"
-        parent_code: "{omop_vocabulary_id}/{omop_concept_code}"
+        parent_codes: "{omop_vocabulary_id}/{omop_concept_code}"
 
 icu/procedureevents:
   start:
@@ -239,11 +239,11 @@ icu/procedureevents:
       proc_datetimeevents:
         description: ["omop_concept_name", "label"] # List of strings are columns to be collated
         itemid: "itemid (omop_source_code)"
-        parent_code: "{omop_vocabulary_id}/{omop_concept_code}"
+        parent_codes: "{omop_vocabulary_id}/{omop_concept_code}"
       proc_itemid:
         description: ["omop_concept_name", "label"] # List of strings are columns to be collated
         itemid: "itemid (omop_source_code)"
-        parent_code: "{omop_vocabulary_id}/{omop_concept_code}"
+        parent_codes: "{omop_vocabulary_id}/{omop_concept_code}"
   end:
     code:
       - PROCEDURE
@@ -257,11 +257,11 @@ icu/procedureevents:
       proc_datetimeevents:
         description: ["omop_concept_name", "label"] # List of strings are columns to be collated
         itemid: "itemid (omop_source_code)"
-        parent_code: "{omop_vocabulary_id}/{omop_concept_code}"
+        parent_codes: "{omop_vocabulary_id}/{omop_concept_code}"
       proc_itemid:
         description: ["omop_concept_name", "label"] # List of strings are columns to be collated
         itemid: "itemid (omop_source_code)"
-        parent_code: "{omop_vocabulary_id}/{omop_concept_code}"
+        parent_codes: "{omop_vocabulary_id}/{omop_concept_code}"
 
 icu/inputevents:
   input_start:
@@ -283,7 +283,7 @@ icu/inputevents:
       #inputevents_to_rxnorm:
       #  description: ["omop_concept_name", "label"] # List of strings are columns to be collated
       #  itemid: "itemid (omop_source_code)"
-      #  parent_code: "{omop_vocabulary_id}/{omop_concept_code}"
+      #  parent_codes: "{omop_vocabulary_id}/{omop_concept_code}"
       #  rateuom: null # A null column means this column is needed in pulling from the metadata.
   input_end:
     code:
@@ -304,7 +304,7 @@ icu/inputevents:
       #  inputevents_to_rxnorm:
       #    description: ["omop_concept_name", "label"] # List of strings are columns to be collated
       #    itemid: "itemid (omop_source_code)"
-      #    parent_code: "{omop_vocabulary_id}/{omop_concept_code}"
+      #    parent_codes: "{omop_vocabulary_id}/{omop_concept_code}"
       #    statusdescription: null
       #    amountuom: null
   patient_weight:
@@ -331,4 +331,4 @@ icu/outputevents:
         description: ["omop_concept_name", "label"] # List of strings are columns to be collated
         itemid: "itemid (omop_source_code)"
         valueuom: "unitname"
-        parent_code: "{omop_vocabulary_id}/{omop_concept_code}"
+        parent_codes: "{omop_vocabulary_id}/{omop_concept_code}"
diff --git a/src/MEDS_transforms/configs/extract.yaml b/src/MEDS_transforms/configs/extract.yaml
index 53f51bb..fb1b05f 100644
--- a/src/MEDS_transforms/configs/extract.yaml
+++ b/src/MEDS_transforms/configs/extract.yaml
@@ -4,6 +4,7 @@ defaults:
       - shard_events
       - split_and_shard_patients
       - merge_to_MEDS_cohort
+      - extract_code_metadata
       # There is no configuration beyond the global "event_conversion_config_fp" for the
       # convert_to_sharded_events stage, so it doesn't have a stage config block here or below.
   - _self_
@@ -57,5 +58,3 @@ stage_configs:
       - "values/sum_sqd"
     do_summarize_over_all_codes: true # This indicates we should include overall, code-independent counts
     process_shard_prefix: "train/" # we only want to summarize over the training set
-  extract_code_metadata:
-    is_metadata: true
diff --git a/src/MEDS_transforms/configs/stage_configs/extract_code_metadata.yaml b/src/MEDS_transforms/configs/stage_configs/extract_code_metadata.yaml
new file mode 100644
index 0000000..09ba0ba
--- /dev/null
+++ b/src/MEDS_transforms/configs/stage_configs/extract_code_metadata.yaml
@@ -0,0 +1,3 @@
+extract_code_metadata:
+  is_metadata: true
+  description_seperator: "\n"
diff --git a/src/MEDS_transforms/extract/extract_code_metadata.py b/src/MEDS_transforms/extract/extract_code_metadata.py
index e789a11..166cc0e 100644
--- a/src/MEDS_transforms/extract/extract_code_metadata.py
+++ b/src/MEDS_transforms/extract/extract_code_metadata.py
@@ -20,6 +20,11 @@ from MEDS_transforms.extract.utils import get_supported_fp
 from MEDS_transforms.mapreduce.mapper import rwlock_wrap
 from MEDS_transforms.utils import stage_init, write_lazyframe
 
+MEDS_METADATA_MANDATORY_TYPES = {
+    "description": pl.Utf8,
+    "parent_codes": pl.List(pl.Utf8),
+}
+
 
 def extract_metadata(
     metadata_df: pl.LazyFrame, event_cfg: dict[str, str | None], allowed_codes: list | None = None
@@ -98,7 +103,8 @@ def extract_metadata(
             ...
         TypeError: Event configuration must be a dictionary. Got: <class 'list'> ['foo'].
 
-    You can also manipulate the columns in more complex ways when assigning metadata from the input source.
+    You can also manipulate the columns in more complex ways when assigning metadata from the input source,
+    and mandatory MEDS metadata columns will be cast to the correct types:
         >>> raw_metadata = pl.DataFrame({
         ...     "code": ["A", "A", "C", "D"],
         ...     "code_modifier": ["1", "1", "2", "3"],
@@ -109,8 +115,8 @@ def extract_metadata(
         >>> event_cfg = {
         ...     "code": ["FOO", "col(code)", "col(code_modifier)"],
         ...     "_metadata": {
-        ...         "desc": ["special_title", "title"],
-        ...         "parent_code": [
+        ...         "description": ["special_title", "title"],
+        ...         "parent_codes": [
         ...             {"OUT_VAL/{code_modifier}/2": {"code_modifier_2": "2"}},
         ...             {"OUT_VAL_for_3/{code_modifier}": {"code_modifier_2": "3"}},
         ...             {
@@ -122,16 +128,16 @@ def extract_metadata(
         ... }
         >>> extract_metadata(raw_metadata, event_cfg)
         shape: (4, 3)
-        ┌───────────┬───────┬─────────────────┐
-        │ code      ┆ desc  ┆ parent_code     │
-        │ ---       ┆ ---   ┆ ---             │
-        │ str       ┆ str   ┆ str             │
-        ╞═══════════╪═══════╪═════════════════╡
-        │ FOO//A//1 ┆ used  ┆ null            │
-        │ FOO//A//1 ┆ A-1-2 ┆ OUT_VAL/1/2     │
-        │ FOO//C//2 ┆ C-2-3 ┆ OUT_VAL_for_3/2 │
-        │ FOO//D//3 ┆ null  ┆ expanded form   │
-        └───────────┴───────┴─────────────────┘
+        ┌───────────┬─────────────┬─────────────────────┐
+        │ code      ┆ description ┆ parent_codes        │
+        │ ---       ┆ ---         ┆ ---                 │
+        │ str       ┆ str         ┆ list[str]           │
+        ╞═══════════╪═════════════╪═════════════════════╡
+        │ FOO//A//1 ┆ used        ┆ null                │
+        │ FOO//A//1 ┆ A-1-2       ┆ ["OUT_VAL/1/2"]     │
+        │ FOO//C//2 ┆ C-2-3       ┆ ["OUT_VAL_for_3/2"] │
+        │ FOO//D//3 ┆ null        ┆ ["expanded form"]   │
+        └───────────┴─────────────┴─────────────────────┘
     """
     event_cfg = copy.deepcopy(event_cfg)
 
@@ -176,6 +182,14 @@ def extract_metadata(
 
     metadata_df = metadata_df.filter(~pl.all_horizontal(*[pl.col(c).is_null() for c in final_cols]))
 
+    for mandatory_col, mandatory_type in MEDS_METADATA_MANDATORY_TYPES.items():
+        if mandatory_col not in final_cols:
+            continue
+
+        if metadata_df.schema[mandatory_col] is not mandatory_type:
+            logger.warning(f"Metadata column '{mandatory_col}' must be of type {mandatory_type}. Casting.")
+            metadata_df = metadata_df.with_columns(pl.col(mandatory_col).cast(mandatory_type), strict=False)
+
     return metadata_df.unique(maintain_order=True).select("code", *final_cols)
 
 
@@ -312,7 +326,33 @@ def get_events_and_metadata_by_metadata_fp(event_configs: dict | DictConfig) ->
 
 @hydra.main(version_base=None, config_path=str(CONFIG_YAML.parent), config_name=CONFIG_YAML.stem)
 def main(cfg: DictConfig):
-    """TODO."""
+    """Extracts any dataset-specific metadata and adds it to any existing code metadata file.
+
+    This script can extract arbitrary, code-linked metadata columns from input mappings and add them to the
+    `metadata/codes.parquet` file. The metadata columns are extracted from the raw metadata files using a
+    parsing DSL that is specified in the `event_conversion_config_fp` file. See `parser.py` for more details
+    on this DSL.
+
+    Note that there are two sentinel columns in the output metadata that have certain mandates for MEDS
+    compliance: The `description` column and the `parent_codes` column. The `description` column must be a
+    string, and if there are multiple matches in the extracted metadata for a code, in this script they will
+    be concatenated into a single string with the `description_separator` string. The `parent_codes` column
+    must be a list of strings, each formatted as an OMOP vocabulary name, followed by a "/", followed by the
+    OMOP concept code. This column is used to link codes to their parent codes in the OMOP vocabulary.
+
+    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. It cannot be overwritten
+    directly on the command line, but can be overwritten implicitly by overwriting components of the
+    `stage_configs.extract_code_metadata` key.
+
+    Args:
+        stage_configs.extract_code_metadata.description_separator: If there are multiple metadata matches for
+            a row, this string will be used as a separator to join the matches for the sentinel
+            `"description"` column into a single string in the output metadata, per compliance with the MEDS
+            schema.
+    """
 
     stage_input_dir, partial_metadata_dir, _, _ = stage_init(cfg)
     raw_input_dir = Path(cfg.input_dir)
@@ -384,9 +424,15 @@ def main(cfg: DictConfig):
     logger.info(f"Collected metadata for {n_unique_obs} unique codes among {n_rows} total observations.")
 
     if n_unique_obs != n_rows:
-        reduced = reduced.group_by(join_cols).agg(*(pl.col(c) for c in metadata_cols)).collect()
-    else:
-        reduced = reduced.collect()
+        aggs = {c: pl.col(c) for c in metadata_cols if c not in MEDS_METADATA_MANDATORY_TYPES}
+        if "description" in metadata_cols:
+            aggs["description"] = pl.col("description").list.join(cfg.stage_cfg.description_separator)
+        if "parent_codes" in metadata_cols:
+            aggs["parent_codes"] = pl.col("parent_codes").explode().implode()
+
+        reduced = reduced.group_by(join_cols).agg(*(pl.col(c) for c in metadata_cols))
+
+    reduced = reduced.collect()
 
     metadata_input_dir = Path(cfg.stage_cfg.metadata_input_dir)
     old_metadata_fp = metadata_input_dir / "codes.parquet"


Test patch:
diff --git a/tests/test_extract.py b/tests/test_extract.py
index bf0942f..f30f974 100644
--- a/tests/test_extract.py
+++ b/tests/test_extract.py
@@ -121,7 +121,7 @@ admit_vitals:
     _metadata:
       input_metadata:
         description: {"title": {"lab_code": "HR"}}
-        parent_code: {"LOINC/{loinc}": {"lab_code": "HR"}}
+        parent_codes: {"LOINC/{loinc}": {"lab_code": "HR"}}
   temp:
     code: TEMP
     timestamp: col(vitals_date)
@@ -130,7 +130,7 @@ admit_vitals:
     _metadata:
       input_metadata:
         description: {"title": {"lab_code": "temp"}}
-        parent_code: {"LOINC/{loinc}": {"lab_code": "temp"}}
+        parent_codes: {"LOINC/{loinc}": {"lab_code": "temp"}}
 """
 
 # Test data (expected outputs) -- ALL OF THIS MAY CHANGE IF THE SEED OR DATA CHANGES
@@ -266,7 +266,7 @@ TEMP,12,4,12,1181.4999999999998,116373.38999999998
 """
 
 MEDS_OUTPUT_CODE_METADATA_FILE_WITH_DESC = """
-code,code/n_occurrences,code/n_patients,values/n_occurrences,values/sum,values/sum_sqd,description,parent_code
+code,code/n_occurrences,code/n_patients,values/n_occurrences,values/sum,values/sum_sqd,description,parent_codes
 ,44,4,28,3198.8389005974336,382968.28937288234,,
 ADMISSION//CARDIAC,2,2,0,,,,
 ADMISSION//ORTHOPEDIC,1,1,0,,,,
@@ -592,8 +592,13 @@ def test_extraction():
             pl.col("values/n_occurrences").cast(pl.UInt8),
             pl.col("values/sum").cast(pl.Float32).fill_null(0),
             pl.col("values/sum_sqd").cast(pl.Float32).fill_null(0),
+            pl.col("parent_codes").cast(pl.List(pl.Utf8)),
         )
 
+        # We collapse the list type as it throws an error in the assert_df_equal otherwise
+        got_df = got_df.with_columns(pl.col("parent_codes").list.join("||"))
+        want_df = want_df.with_columns(pl.col("parent_codes").list.join("||"))
+
         assert_df_equal(
             want=want_df,
             got=got_df,