Repository: mmcdermott/MEDS_transforms

Issue:
The default extraction ETL should likely not include an `aggregate_code_metadata.py` stage, unless anyone thinks it would be almost universally useful.
This means columns like `code/n_occurrences`, `value/sum`, etc. would not be computed during aggregation. Code metadata (e.g., `description`, `parent_codes`, etc.) would still be included by default. 

Does anyone see any reason why this aggregation stage should be included by default during extraction to MEDS (the stage will still be usable of course during pre-processing pipelines).

Tagging @EthanSteinberg, @Oufattole, @prenc, @prockenschaub, @tompollard for inputs.

Per discussion below, there are two planned tasks:
- [x] remove `aggregate_code_metadata.py` from the default ETL
- [ ] Make `finalize_MEDS_metadata.py` ensure that there are rows for all unique codes in the dataset in `metadata/codes.parquet`, even if the rows for some or all codes are all null. (this will be relegated to #117 
- [x] Add separate, standalone integration tests for `aggregate_code_metadata.py` as opposed to having it tested through the extraction ETL.

Discussion:

Patch:
diff --git a/src/MEDS_transforms/aggregate_code_metadata.py b/src/MEDS_transforms/aggregate_code_metadata.py
index 9290926..bcc056a 100755
--- a/src/MEDS_transforms/aggregate_code_metadata.py
+++ b/src/MEDS_transforms/aggregate_code_metadata.py
@@ -682,6 +682,15 @@ def run_map_reduce(cfg: DictConfig):
         cs.numeric().shrink_dtype().name.keep()
     )
 
+    old_metadata_fp = Path(cfg.stage_cfg.metadata_input_dir) / "codes.parquet"
+    join_cols = ["code", *cfg.get("code_modifier_cols", [])]
+
+    if old_metadata_fp.exists():
+        logger.info(f"Joining to existing code metadata at {str(old_metadata_fp.resolve())}")
+        existing = pl.scan_parquet(old_metadata_fp)
+        existing = existing.drop(*[c for c in existing.columns if c in set(reduced.columns) - set(join_cols)])
+        reduced = reduced.join(existing, on=join_cols, how="left", coalesce=True)
+
     write_lazyframe(reduced, reducer_fp)
     logger.info(f"Finished reduction in {datetime.now() - start}")
 
diff --git a/src/MEDS_transforms/configs/extract.yaml b/src/MEDS_transforms/configs/extract.yaml
index b0f2d50..ae3bf07 100644
--- a/src/MEDS_transforms/configs/extract.yaml
+++ b/src/MEDS_transforms/configs/extract.yaml
@@ -27,8 +27,6 @@ description: |-
 
 # The event conversion configuration file is used throughout the pipeline to define the events to extract.
 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}/metadata/.shards.json"
 
@@ -37,7 +35,6 @@ stages:
   - split_and_shard_patients
   - convert_to_sharded_events
   - merge_to_MEDS_cohort
-  - aggregate_code_metadata
   - extract_code_metadata
   - finalize_MEDS_metadata
   - finalize_MEDS_data
@@ -45,19 +42,3 @@ stages:
 stage_configs:
   shard_events:
     data_input_dir: "${input_dir}"
-  aggregate_code_metadata:
-    description: |-
-      This stage collects some descriptive metadata about the codes in the cohort.
-
-      Args:
-        stage_cfg.aggregations: The aggregations to compute over the codes.
-          Defaults to counts of code occurrences, counts of patients with the code, and counts of value
-          occurrences per code, as well as the sum and sum of squares of values (for use in computing means
-          and variances).
-    aggregations:
-      - "code/n_occurrences"
-      - "code/n_patients"
-      - "values/n_occurrences"
-      - "values/sum"
-      - "values/sum_sqd"
-    do_summarize_over_all_codes: true # This indicates we should include overall, code-independent counts


Test patch:
diff --git a/tests/test_aggregate_code_metadata.py b/tests/test_aggregate_code_metadata.py
new file mode 100644
index 0000000..2cdc02a
--- /dev/null
+++ b/tests/test_aggregate_code_metadata.py
@@ -0,0 +1,55 @@
+"""Tests the aggregate_code_metadata.py script.
+
+Set the bash env variable `DO_USE_LOCAL_SCRIPTS=1` to use the local py files, rather than the installed
+scripts.
+"""
+
+from .transform_tester_base import (
+    AGGREGATE_CODE_METADATA_SCRIPT,
+    parse_code_metadata_csv,
+    single_stage_transform_tester,
+)
+
+WANT_OUTPUT_CODE_METADATA_FILE = """
+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,0,0,,
+ADMISSION//ORTHOPEDIC,1,1,0,0,0,,
+ADMISSION//PULMONARY,1,1,0,0,0,,
+DISCHARGE,4,4,0,0,0,,
+DOB,4,4,0,0,0,,
+EYE_COLOR//BLUE,1,1,0,0,0,"Blue Eyes. Less common than brown.",
+EYE_COLOR//BROWN,1,1,0,0,0,"Brown Eyes. The most common eye color.",
+EYE_COLOR//HAZEL,2,2,0,0,0,"Hazel eyes. These are uncommon",
+HEIGHT,4,4,4,656.8389005974336,108056.12937288235,,
+HR,12,4,12,1360.5000000000002,158538.77,"Heart Rate",LOINC/8867-4
+TEMP,12,4,12,1181.4999999999998,116373.38999999998,"Body Temperature",LOINC/8310-5
+"""
+
+MEDS_CODE_METADATA_FILE = """
+code,description,parent_codes
+EYE_COLOR//BLUE,"Blue Eyes. Less common than brown.",
+EYE_COLOR//BROWN,"Brown Eyes. The most common eye color.",
+EYE_COLOR//HAZEL,"Hazel eyes. These are uncommon",
+HR,"Heart Rate",LOINC/8867-4
+TEMP,"Body Temperature",LOINC/8310-5
+"""
+
+AGGREGATIONS = [
+    "code/n_occurrences",
+    "code/n_patients",
+    "values/n_occurrences",
+    "values/sum",
+    "values/sum_sqd",
+]
+
+
+def test_aggregate_code_metadata():
+    single_stage_transform_tester(
+        transform_script=AGGREGATE_CODE_METADATA_SCRIPT,
+        stage_name="aggregate_code_metadata",
+        transform_stage_kwargs={"aggregations": AGGREGATIONS, "do_summarize_over_all_codes": True},
+        want_outputs=parse_code_metadata_csv(WANT_OUTPUT_CODE_METADATA_FILE),
+        code_metadata=MEDS_CODE_METADATA_FILE,
+        do_use_config_yaml=True,
+    )
diff --git a/tests/test_extract.py b/tests/test_extract.py
index b59ae02..d8a3c3e 100644
--- a/tests/test_extract.py
+++ b/tests/test_extract.py
@@ -18,7 +18,6 @@ if os.environ.get("DO_USE_LOCAL_SCRIPTS", "0") == "1":
     SPLIT_AND_SHARD_SCRIPT = extraction_root / "split_and_shard_patients.py"
     CONVERT_TO_SHARDED_EVENTS_SCRIPT = extraction_root / "convert_to_sharded_events.py"
     MERGE_TO_MEDS_COHORT_SCRIPT = extraction_root / "merge_to_MEDS_cohort.py"
-    AGGREGATE_CODE_METADATA_SCRIPT = code_root / "aggregate_code_metadata.py"
     EXTRACT_CODE_METADATA_SCRIPT = extraction_root / "extract_code_metadata.py"
     FINALIZE_DATA_SCRIPT = extraction_root / "finalize_MEDS_data.py"
     FINALIZE_METADATA_SCRIPT = extraction_root / "finalize_MEDS_metadata.py"
@@ -27,7 +26,6 @@ else:
     SPLIT_AND_SHARD_SCRIPT = "MEDS_extract-split_and_shard_patients"
     CONVERT_TO_SHARDED_EVENTS_SCRIPT = "MEDS_extract-convert_to_sharded_events"
     MERGE_TO_MEDS_COHORT_SCRIPT = "MEDS_extract-merge_to_MEDS_cohort"
-    AGGREGATE_CODE_METADATA_SCRIPT = "MEDS_transform-aggregate_code_metadata"
     EXTRACT_CODE_METADATA_SCRIPT = "MEDS_extract-extract_code_metadata"
     FINALIZE_DATA_SCRIPT = "MEDS_extract-finalize_MEDS_data"
     FINALIZE_METADATA_SCRIPT = "MEDS_extract-finalize_MEDS_metadata"
@@ -262,35 +260,12 @@ patient_id,time,code,numeric_value
 """
 
 MEDS_OUTPUT_CODE_METADATA_FILE = """
-code,code/n_occurrences,code/n_patients,values/n_occurrences,values/sum,values/sum_sqd
-,44,4,28,3198.8389005974336,382968.28937288234
-ADMISSION//CARDIAC,2,2,0,,
-ADMISSION//ORTHOPEDIC,1,1,0,,
-ADMISSION//PULMONARY,1,1,0,,
-DISCHARGE,4,4,0,,
-DOB,4,4,0,,
-EYE_COLOR//BLUE,1,1,0,,
-EYE_COLOR//BROWN,1,1,0,,
-EYE_COLOR//HAZEL,2,2,0,,
-HEIGHT,4,4,4,656.8389005974336,108056.12937288235
-HR,12,4,12,1360.5000000000002,158538.77
-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_codes
-,44,4,28,3198.8389005974336,382968.28937288234,,
-ADMISSION//CARDIAC,2,2,0,,,,
-ADMISSION//ORTHOPEDIC,1,1,0,,,,
-ADMISSION//PULMONARY,1,1,0,,,,
-DISCHARGE,4,4,0,,,,
-DOB,4,4,0,,,,
-EYE_COLOR//BLUE,1,1,0,,,"Blue Eyes. Less common than brown.",
-EYE_COLOR//BROWN,1,1,0,,,"Brown Eyes. The most common eye color.",
-EYE_COLOR//HAZEL,2,2,0,,,"Hazel eyes. These are uncommon",
-HEIGHT,4,4,4,656.8389005974336,108056.12937288235,,
-HR,12,4,12,1360.5000000000002,158538.77,"Heart Rate",LOINC/8867-4
-TEMP,12,4,12,1181.4999999999998,116373.38999999998,"Body Temperature",LOINC/8310-5
+code,description,parent_codes
+EYE_COLOR//BLUE,"Blue Eyes. Less common than brown.",
+EYE_COLOR//BROWN,"Brown Eyes. The most common eye color.",
+EYE_COLOR//HAZEL,"Hazel eyes. These are uncommon",
+HR,"Heart Rate",LOINC/8867-4
+TEMP,"Body Temperature",LOINC/8310-5
 """
 
 MEDS_OUTPUT_DATASET_METADATA_JSON = {
@@ -549,41 +524,6 @@ def test_extraction():
             print(f"stdout:\n{full_stdout}")
             raise e
 
-        # Stage 5: Aggregate preliminary code metadata
-        stderr, stdout = run_command(
-            AGGREGATE_CODE_METADATA_SCRIPT,
-            extraction_config_kwargs,
-            "aggregate_code_metadata",
-            config_name="extract",
-        )
-        all_stderrs.append(stderr)
-        all_stdouts.append(stdout)
-
-        full_stderr = "\n".join(all_stderrs)
-        full_stdout = "\n".join(all_stdouts)
-
-        output_file = MEDS_cohort_dir / "aggregate_code_metadata" / "codes.parquet"
-        assert output_file.is_file(), f"Expected {output_file} to exist: stderr:\n{stderr}\nstdout:\n{stdout}"
-
-        got_df = pl.read_parquet(output_file, glob=False)
-
-        want_df = pl.read_csv(source=StringIO(MEDS_OUTPUT_CODE_METADATA_FILE)).with_columns(
-            pl.col("code"),
-            pl.col("code/n_occurrences").cast(pl.UInt8),
-            pl.col("code/n_patients").cast(pl.UInt8),
-            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),
-        )
-
-        assert_df_equal(
-            want=want_df,
-            got=got_df,
-            msg="Code metadata differs!",
-            check_column_order=False,
-            check_row_order=False,
-        )
-
         # Stage 6: Extract code metadata
         stderr, stdout = run_command(
             EXTRACT_CODE_METADATA_SCRIPT,
@@ -601,13 +541,8 @@ def test_extraction():
 
         got_df = pl.read_parquet(output_file, glob=False)
 
-        want_df = pl.read_csv(source=StringIO(MEDS_OUTPUT_CODE_METADATA_FILE_WITH_DESC)).with_columns(
+        want_df = pl.read_csv(source=StringIO(MEDS_OUTPUT_CODE_METADATA_FILE)).with_columns(
             pl.col("code"),
-            pl.col("code/n_occurrences").cast(pl.UInt8),
-            pl.col("code/n_patients").cast(pl.UInt8),
-            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)),
         )
 
@@ -689,13 +624,8 @@ def test_extraction():
 
         got_df = pl.read_parquet(output_file, glob=False, use_pyarrow=True)
 
-        want_df = pl.read_csv(source=StringIO(MEDS_OUTPUT_CODE_METADATA_FILE_WITH_DESC)).with_columns(
+        want_df = pl.read_csv(source=StringIO(MEDS_OUTPUT_CODE_METADATA_FILE)).with_columns(
             pl.col("code"),
-            pl.col("code/n_occurrences").cast(pl.UInt8),
-            pl.col("code/n_patients").cast(pl.UInt8),
-            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)),
         )
 
diff --git a/tests/test_fit_vocabulary_indices.py b/tests/test_fit_vocabulary_indices.py
index 52d2484..312e65f 100644
--- a/tests/test_fit_vocabulary_indices.py
+++ b/tests/test_fit_vocabulary_indices.py
@@ -4,18 +4,15 @@ Set the bash env variable `DO_USE_LOCAL_SCRIPTS=1` to use the local py files, ra
 scripts.
 """
 
-from io import StringIO
-
-import polars as pl
 
 from .transform_tester_base import (
     FIT_VOCABULARY_INDICES_SCRIPT,
-    MEDS_CODE_METADATA_SCHEMA,
+    parse_code_metadata_csv,
     single_stage_transform_tester,
 )
 
 WANT_CSV = """
-code,code/n_occurrences,code/n_patients,values/n_occurrences,values/sum,values/sum_sqd,description,parent_code,code/vocab_index
+code,code/n_occurrences,code/n_patients,values/n_occurrences,values/sum,values/sum_sqd,description,parent_codes,code/vocab_index
 ,44,4,28,3198.8389005974336,382968.28937288234,,,1
 ADMISSION//CARDIAC,2,2,0,,,,,2
 ADMISSION//ORTHOPEDIC,1,1,0,,,,,3
@@ -30,18 +27,11 @@ HR,12,4,12,1360.5000000000002,158538.77,"Heart Rate",LOINC/8867-4,11
 TEMP,12,4,12,1181.4999999999998,116373.38999999998,"Body Temperature",LOINC/8310-5,12
 """
 
-WANT_SCHEMA = {
-    **MEDS_CODE_METADATA_SCHEMA,
-    "code/vocab_index": pl.UInt8,
-}
-
-WANT_DF = pl.read_csv(StringIO(WANT_CSV), schema=WANT_SCHEMA)
-
 
 def test_fit_vocabulary_indices_with_default_stage_config():
     single_stage_transform_tester(
         transform_script=FIT_VOCABULARY_INDICES_SCRIPT,
         stage_name="fit_vocabulary_indices",
         transform_stage_kwargs=None,
-        want_outputs=WANT_DF,
+        want_outputs=parse_code_metadata_csv(WANT_CSV),
     )
diff --git a/tests/transform_tester_base.py b/tests/transform_tester_base.py
index 7fa5ba9..2efeb02 100644
--- a/tests/transform_tester_base.py
+++ b/tests/transform_tester_base.py
@@ -26,6 +26,7 @@ filters_root = code_root / "filters"
 
 if os.environ.get("DO_USE_LOCAL_SCRIPTS", "0") == "1":
     # Root Source
+    AGGREGATE_CODE_METADATA_SCRIPT = code_root / "aggregate_code_metadata.py"
     FIT_VOCABULARY_INDICES_SCRIPT = code_root / "fit_vocabulary_indices.py"
     RESHARD_TO_SPLIT_SCRIPT = code_root / "reshard_to_split.py"
 
@@ -42,6 +43,7 @@ if os.environ.get("DO_USE_LOCAL_SCRIPTS", "0") == "1":
     TOKENIZATION_SCRIPT = transforms_root / "tokenization.py"
 else:
     # Root Source
+    AGGREGATE_CODE_METADATA_SCRIPT = "MEDS_transform-aggregate_code_metadata"
     FIT_VOCABULARY_INDICES_SCRIPT = "MEDS_transform-fit_vocabulary_indices"
     RESHARD_TO_SPLIT_SCRIPT = "MEDS_transform-reshard_to_split"
 
@@ -161,7 +163,7 @@ MEDS_SHARDS = parse_meds_csvs(
 
 
 MEDS_CODE_METADATA_CSV = """
-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,,,,
@@ -184,7 +186,7 @@ MEDS_CODE_METADATA_SCHEMA = {
     "values/sum": pl.Float32,
     "values/sum_sqd": pl.Float32,
     "description": pl.Utf8,
-    "parent_code": pl.Utf8,
+    "parent_codes": pl.Utf8,
     "code/vocab_index": pl.UInt8,
 }
 
@@ -192,7 +194,10 @@ MEDS_CODE_METADATA_SCHEMA = {
 def parse_code_metadata_csv(csv_str: str) -> pl.DataFrame:
     cols = csv_str.strip().split("\n")[0].split(",")
     schema = {col: dt for col, dt in MEDS_CODE_METADATA_SCHEMA.items() if col in cols}
-    return pl.read_csv(StringIO(csv_str), schema=schema)
+    df = pl.read_csv(StringIO(csv_str), schema=schema)
+    if "parent_codes" in cols:
+        df = df.with_columns(pl.col("parent_codes").cast(pl.List(pl.Utf8)))
+    return df
 
 
 MEDS_CODE_METADATA = parse_code_metadata_csv(MEDS_CODE_METADATA_CSV)