Repository: mmcdermott/MEDS_transforms

Issue:
Aggregation integration test should cover all integrations

Discussion:

Patch:
diff --git a/src/MEDS_transforms/aggregate_code_metadata.py b/src/MEDS_transforms/aggregate_code_metadata.py
index bcc056a..13e9b34 100755
--- a/src/MEDS_transforms/aggregate_code_metadata.py
+++ b/src/MEDS_transforms/aggregate_code_metadata.py
@@ -406,7 +406,9 @@ def mapper_fntr(
         │ C    ┆ 1         ┆ 81.25          ┆ 5.0        ┆ 7.5        │
         │ D    ┆ null      ┆ 0.0            ┆ null       ┆ null       │
         └──────┴───────────┴────────────────┴────────────┴────────────┘
-        >>> stage_cfg = DictConfig({"aggregations": ["values/quantiles"]})
+        >>> stage_cfg = DictConfig({
+        ...     "aggregations": [{"name": "values/quantiles", "quantiles": [0.25, 0.5, 0.75]}]
+        ... })
         >>> mapper = mapper_fntr(stage_cfg, code_modifiers)
         >>> mapper(df.lazy()).collect().select("code", "modifier1", pl.col("values/quantiles"))
         shape: (5, 3)
@@ -421,6 +423,25 @@ def mapper_fntr(
         │ C    ┆ 1         ┆ [5.0, 7.5]       │
         │ D    ┆ null      ┆ []               │
         └──────┴───────────┴──────────────────┘
+        >>> stage_cfg = DictConfig({
+        ...     "aggregations": [{"name": "values/quantiles", "quantiles": [0.25, 0.5, 0.75]}],
+        ...     "do_summarize_over_all_codes": True,
+        ... })
+        >>> mapper = mapper_fntr(stage_cfg, code_modifiers)
+        >>> mapper(df.lazy()).collect().select("code", "modifier1", pl.col("values/quantiles"))
+        shape: (6, 3)
+        ┌──────┬───────────┬───────────────────┐
+        │ code ┆ modifier1 ┆ values/quantiles  │
+        │ ---  ┆ ---       ┆ ---               │
+        │ str  ┆ i64       ┆ list[f64]         │
+        ╞══════╪═══════════╪═══════════════════╡
+        │ null ┆ null      ┆ [1.1, 2.0, … 7.5] │
+        │ A    ┆ 1         ┆ [1.1, 1.1]        │
+        │ A    ┆ 2         ┆ [6.0]             │
+        │ B    ┆ 2         ┆ [2.0, 4.0]        │
+        │ C    ┆ 1         ┆ [5.0, 7.5]        │
+        │ D    ┆ null      ┆ []                │
+        └──────┴───────────┴───────────────────┘
     """
 
     code_key_columns = validate_args_and_get_code_cols(stage_cfg, code_modifiers)
@@ -435,7 +456,12 @@ def mapper_fntr(
         return df.group_by(code_key_columns).agg(**agg_operations).sort(code_key_columns)
 
     def all_patients_mapper(df: pl.LazyFrame) -> pl.LazyFrame:
-        return df.select(**agg_operations)
+        local_agg_operations = agg_operations.copy()
+        if METADATA_FN.VALUES_QUANTILES in agg_operations:
+            local_agg_operations[METADATA_FN.VALUES_QUANTILES] = agg_operations[
+                METADATA_FN.VALUES_QUANTILES
+            ].implode()
+        return df.select(**local_agg_operations)
 
     if stage_cfg.get("do_summarize_over_all_codes", False):
 
@@ -443,7 +469,7 @@ def mapper_fntr(
             by_code = by_code_mapper(df)
             all_patients = all_patients_mapper(df)
             return pl.concat([all_patients, by_code], how="diagonal_relaxed").select(
-                *code_key_columns, *aggregations
+                *code_key_columns, *agg_operations.keys()
             )
 
     else:


Test patch:
diff --git a/tests/test_aggregate_code_metadata.py b/tests/test_aggregate_code_metadata.py
index 2cdc02a..2a8f78c 100644
--- a/tests/test_aggregate_code_metadata.py
+++ b/tests/test_aggregate_code_metadata.py
@@ -4,28 +4,154 @@ Set the bash env variable `DO_USE_LOCAL_SCRIPTS=1` to use the local py files, ra
 scripts.
 """
 
+import polars as pl
+
 from .transform_tester_base import (
     AGGREGATE_CODE_METADATA_SCRIPT,
-    parse_code_metadata_csv,
+    MEDS_CODE_METADATA_SCHEMA,
     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
+code,code/n_occurrences,code/n_patients,values/n_occurrences,values/n_patients,values/sum,values/sum_sqd,values/n_ints,values/min,values/max,description,parent_codes
+,44,4,28,4,3198.8389005974336,382968.28937288234,6,86.0,175.271118,,
+ADMISSION//CARDIAC,2,2,0,0,0,0,0,,,,
+ADMISSION//ORTHOPEDIC,1,1,0,0,0,0,0,,,,
+ADMISSION//PULMONARY,1,1,0,0,0,0,0,,,,
+DISCHARGE,4,4,0,0,0,0,0,,,,
+DOB,4,4,0,0,0,0,0,,,,
+EYE_COLOR//BLUE,1,1,0,0,0,0,0,,,"Blue Eyes. Less common than brown.",
+EYE_COLOR//BROWN,1,1,0,0,0,0,0,,,"Brown Eyes. The most common eye color.",
+EYE_COLOR//HAZEL,2,2,0,0,0,0,0,,,"Hazel eyes. These are uncommon",
+HEIGHT,4,4,4,4,656.8389005974336,108056.12937288235,0,156.485596,175.271118,,
+HR,12,4,12,4,1360.5000000000002,158538.77,2,86.0,170.199997,"Heart Rate",LOINC/8867-4
+TEMP,12,4,12,4,1181.4999999999998,116373.38999999998,4,95.5,100.400002,"Body Temperature",LOINC/8310-5
 """
 
+WANT_OUTPUT_CODE_METADATA_FILE = pl.DataFrame(
+    {
+        "code": [
+            None,
+            "ADMISSION//CARDIAC",
+            "ADMISSION//ORTHOPEDIC",
+            "ADMISSION//PULMONARY",
+            "DISCHARGE",
+            "DOB",
+            "EYE_COLOR//BLUE",
+            "EYE_COLOR//BROWN",
+            "EYE_COLOR//HAZEL",
+            "HEIGHT",
+            "HR",
+            "TEMP",
+        ],
+        "code/n_occurrences": [44, 2, 1, 1, 4, 4, 1, 1, 2, 4, 12, 12],
+        "code/n_patients": [4, 2, 1, 1, 4, 4, 1, 1, 2, 4, 4, 4],
+        "values/n_occurrences": [28, 0, 0, 0, 0, 0, 0, 0, 0, 4, 12, 12],
+        "values/n_patients": [4, 0, 0, 0, 0, 0, 0, 0, 0, 4, 4, 4],
+        "values/sum": [
+            3198.8389005974336,
+            0,
+            0,
+            0,
+            0,
+            0,
+            0,
+            0,
+            0,
+            656.8389005974336,
+            1360.5000000000002,
+            1181.4999999999998,
+        ],
+        "values/sum_sqd": [
+            382968.28937288234,
+            0,
+            0,
+            0,
+            0,
+            0,
+            0,
+            0,
+            0,
+            108056.12937288235,
+            158538.77,
+            116373.38999999998,
+        ],
+        "values/n_ints": [6, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 4],
+        "values/min": [86.0, None, None, None, None, None, None, None, None, 156.485596, 86.0, 95.5],
+        "values/max": [
+            175.271118,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            175.271118,
+            170.199997,
+            100.400002,
+        ],
+        "values/quantiles": [
+            {"values/quantile/0.25": 99.9, "values/quantile/0.5": 105.1, "values/quantile/0.75": 113.4},
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            {
+                "values/quantile/0.25": 160.395311,
+                "values/quantile/0.5": 164.686884,
+                "values/quantile/0.75": 164.686884,
+            },
+            {"values/quantile/0.25": 107.5, "values/quantile/0.5": 112.5, "values/quantile/0.75": 113.4},
+            {"values/quantile/0.25": 96.2, "values/quantile/0.5": 99.9, "values/quantile/0.75": 100.0},
+        ],
+        "description": [
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            "Blue Eyes. Less common than brown.",
+            "Brown Eyes. The most common eye color.",
+            "Hazel eyes. These are uncommon",
+            None,
+            "Heart Rate",
+            "Body Temperature",
+        ],
+        "parent_codes": [
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            ["LOINC/8867-4"],
+            ["LOINC/8310-5"],
+        ],
+    },
+    schema={
+        **{k: v for k, v in MEDS_CODE_METADATA_SCHEMA.items() if k != "code/vocab_index"},
+        "parent_codes": pl.List(pl.String),
+        "values/quantiles": pl.Struct(
+            {
+                "values/quantile/0.25": pl.Float32,
+                "values/quantile/0.5": pl.Float32,
+                "values/quantile/0.75": pl.Float32,
+            }
+        ),
+    },
+)
+
 MEDS_CODE_METADATA_FILE = """
 code,description,parent_codes
 EYE_COLOR//BLUE,"Blue Eyes. Less common than brown.",
@@ -39,8 +165,13 @@ AGGREGATIONS = [
     "code/n_occurrences",
     "code/n_patients",
     "values/n_occurrences",
+    "values/n_patients",
     "values/sum",
     "values/sum_sqd",
+    "values/n_ints",
+    "values/min",
+    "values/max",
+    {"name": "values/quantiles", "quantiles": [0.25, 0.5, 0.75]},
 ]
 
 
@@ -49,7 +180,7 @@ def test_aggregate_code_metadata():
         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),
+        want_outputs=WANT_OUTPUT_CODE_METADATA_FILE,
         code_metadata=MEDS_CODE_METADATA_FILE,
         do_use_config_yaml=True,
     )
diff --git a/tests/transform_tester_base.py b/tests/transform_tester_base.py
index 2efeb02..8dbd637 100644
--- a/tests/transform_tester_base.py
+++ b/tests/transform_tester_base.py
@@ -183,8 +183,12 @@ MEDS_CODE_METADATA_SCHEMA = {
     "code/n_occurrences": pl.UInt8,
     "code/n_patients": pl.UInt8,
     "values/n_occurrences": pl.UInt8,
+    "values/n_patients": pl.UInt8,
     "values/sum": pl.Float32,
     "values/sum_sqd": pl.Float32,
+    "values/n_ints": pl.UInt8,
+    "values/min": pl.Float32,
+    "values/max": pl.Float32,
     "description": pl.Utf8,
     "parent_codes": pl.Utf8,
     "code/vocab_index": pl.UInt8,