Repository: mmcdermott/MEDS_transforms

Issue:
Add match & revise syntax transformations
Suggested syntax / configuration.

Suppose a transformation has a `stage_cfg` that normally takes some arguments. As all `stage_cfg`s are dictionaries, we can do the following:

If (for transformations that support match & revise), the `stage_cfg` has a `_match_and_revise` key, then:
  1. All other `stage_cfg` keys are assumed to be global default values.
  2. An error is thrown if each element of the `_match_and_revise` is not a list of dictionaries with each having at least a `_matcher` key which points to a valid `match_configuration` object (validity criteria yet to be determined).
  3. The transformation will be run for each element of the `_match_and_revise` list filtered to only those measurements that match the `_matcher` key value, using all other keyword arguments of the list element as the `stage_cfg` arguments for that revision run, defaulting to global arguments.

Discussion:

Patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index b188f48..98bc99d 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -126,8 +126,4 @@ repos:
       - id: nbqa-isort
         args: ["--profile=black"]
       - id: nbqa-flake8
-        args:
-          [
-            "--extend-ignore=E203,E402,E501,F401,F841",
-            "--exclude=logs/*,data/*",
-          ]
+        args: ["--extend-ignore=E203,E402,E501,F401,F841", "--exclude=logs/*,data/*"]
diff --git a/eICU_Example/configs/table_preprocessors.yaml b/eICU_Example/configs/table_preprocessors.yaml
index 3faf4aa..a3ad2c3 100644
--- a/eICU_Example/configs/table_preprocessors.yaml
+++ b/eICU_Example/configs/table_preprocessors.yaml
@@ -2,8 +2,7 @@ admissiondx:
   offset_col: "admitdxenteredoffset"
   pseudotime_col: "admitDxEnteredTimestamp"
   output_data_cols: ["admitdxname", "admitdxid"]
-  warning_items:
-    ["How should we use `admitdxtest`?", "How should we use `admitdxpath`?"]
+  warning_items: ["How should we use `admitdxtest`?", "How should we use `admitdxpath`?"]
 
 allergy:
   offset_col: "allergyenteredoffset"
diff --git a/src/MEDS_transforms/extract/convert_to_sharded_events.py b/src/MEDS_transforms/extract/convert_to_sharded_events.py
index 8d6a01e..ee4e9d7 100755
--- a/src/MEDS_transforms/extract/convert_to_sharded_events.py
+++ b/src/MEDS_transforms/extract/convert_to_sharded_events.py
@@ -17,12 +17,7 @@ from omegaconf.listconfig import ListConfig
 from MEDS_transforms.extract import CONFIG_YAML
 from MEDS_transforms.extract.shard_events import META_KEYS
 from MEDS_transforms.mapreduce.mapper import rwlock_wrap
-from MEDS_transforms.utils import (
-    is_col_field,
-    parse_col_field,
-    stage_init,
-    write_lazyframe,
-)
+from MEDS_transforms.utils import is_col_field, parse_col_field, stage_init, write_lazyframe
 
 
 def in_format(fmt: str, ts_name: str) -> pl.Expr:
diff --git a/src/MEDS_transforms/mapreduce/mapper.py b/src/MEDS_transforms/mapreduce/mapper.py
index 6b637f3..61ef75f 100644
--- a/src/MEDS_transforms/mapreduce/mapper.py
+++ b/src/MEDS_transforms/mapreduce/mapper.py
@@ -4,14 +4,16 @@ import inspect
 from collections.abc import Callable, Generator
 from datetime import datetime
 from enum import Enum, auto
-from functools import partial
+from functools import partial, wraps
 from pathlib import Path
 from typing import Any, NotRequired, TypedDict, TypeVar
 
+import hydra
 import polars as pl
 from loguru import logger
-from omegaconf import DictConfig
+from omegaconf import DictConfig, ListConfig
 
+from ..extract.parser import is_matcher, matcher_to_expr
 from ..utils import stage_init, write_lazyframe
 from .utils import rwlock_wrap, shard_iterator
 
@@ -135,17 +137,437 @@ def compute_fn_type(compute_fn: ANY_COMPUTE_FN_T) -> ComputeFnType | None:
 
 
 def identity_fn(df: Any) -> Any:
+    """A "null" compute function that returns the input DataFrame as is.
+
+    Args:
+        df: The input DataFrame.
+
+    Returns:
+        The input DataFrame.
+
+    Examples:
+        >>> df = pl.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
+        >>> (identity_fn(df) == df).select(pl.all_horizontal(pl.all().all())).item()
+        True
+    """
+
     return df
 
 
-def read_and_filter_fntr(patients: list[int], read_fn: Callable[[Path], DF_T]) -> Callable[[Path], DF_T]:
+def read_and_filter_fntr(filter_expr: pl.Expr, read_fn: Callable[[Path], DF_T]) -> Callable[[Path], DF_T]:
+    """Create a function that reads a DataFrame from a file and filters it based on a given expression.
+
+    This is specified as a functor in this way to allow it to modify arbitrary other read functions for use in
+    different mapreduce pipelines.
+
+    Args:
+        filter_expr: The filter expression to apply to the DataFrame.
+        read_fn: The read function to use to read the DataFrame.
+
+    Returns:
+        A function that reads a DataFrame from a file and filters it based on the given expression.
+
+    Examples:
+        >>> dfs = {
+        ...     "df1": pl.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}),
+        ...     "df2": pl.DataFrame({"a": [4, 5, 6], "b": [7, 8, 9]})
+        ... }
+        >>> read_fn = lambda key: dfs[key]
+        >>> fn = read_and_filter_fntr((pl.col("a") % 2) == 0, read_fn)
+        >>> fn("df1")
+        shape: (1, 2)
+        ┌─────┬─────┐
+        │ a   ┆ b   │
+        │ --- ┆ --- │
+        │ i64 ┆ i64 │
+        ╞═════╪═════╡
+        │ 2   ┆ 5   │
+        └─────┴─────┘
+        >>> fn("df2")
+        shape: (2, 2)
+        ┌─────┬─────┐
+        │ a   ┆ b   │
+        │ --- ┆ --- │
+        │ i64 ┆ i64 │
+        ╞═════╪═════╡
+        │ 4   ┆ 7   │
+        │ 6   ┆ 9   │
+        └─────┴─────┘
+        >>> fn = read_and_filter_fntr((pl.col("b") % 2) == 0, read_fn)
+        >>> fn("df1")
+        shape: (2, 2)
+        ┌─────┬─────┐
+        │ a   ┆ b   │
+        │ --- ┆ --- │
+        │ i64 ┆ i64 │
+        ╞═════╪═════╡
+        │ 1   ┆ 4   │
+        │ 3   ┆ 6   │
+        └─────┴─────┘
+        >>> fn("df2")
+        shape: (1, 2)
+        ┌─────┬─────┐
+        │ a   ┆ b   │
+        │ --- ┆ --- │
+        │ i64 ┆ i64 │
+        ╞═════╪═════╡
+        │ 5   ┆ 8   │
+        └─────┴─────┘
+    """
+
     def read_and_filter(in_fp: Path) -> DF_T:
-        df = read_fn(in_fp)
-        return df.filter(pl.col("patient_id").isin(patients))
+        return read_fn(in_fp).filter(filter_expr)
 
     return read_and_filter
 
 
+MATCH_REVISE_KEY = "_match_revise"
+MATCHER_KEY = "_matcher"
+
+
+def is_match_revise(stage_cfg: DictConfig) -> bool:
+    """Check if the stage configuration is in a match and revise format.
+
+    Examples:
+        >>> is_match_revise(DictConfig({"_match_revise": []}))
+        False
+        >>> is_match_revise(DictConfig({"_match_revise": [{"_matcher": {"code": "CODE//TEMP"}}]}))
+        True
+        >>> is_match_revise(DictConfig({"foo": "bar"}))
+        False
+    """
+    return bool(stage_cfg.get(MATCH_REVISE_KEY, False))
+
+
+def validate_match_revise(stage_cfg: DictConfig):
+    """Validate that the stage configuration is in a match and revise format.
+
+    Examples:
+        >>> validate_match_revise(DictConfig({"foo": []}))
+        Traceback (most recent call last):
+            ...
+        ValueError: Stage configuration must contain a _match_revise key
+        >>> validate_match_revise(DictConfig({"_match_revise": "foo"}))
+        Traceback (most recent call last):
+            ...
+        ValueError: Match revise options must be a list, got <class 'str'>
+        >>> validate_match_revise(DictConfig({"_match_revise": [1]}))
+        Traceback (most recent call last):
+            ...
+        ValueError: Match revise config 0 must be a dict, got <class 'int'>
+        >>> validate_match_revise(DictConfig({"_match_revise": [{"_matcher": {"foo": "bar"}}, 1]}))
+        Traceback (most recent call last):
+            ...
+        ValueError: Match revise config 1 must be a dict, got <class 'int'>
+        >>> validate_match_revise(DictConfig({"_match_revise": [{"foo": "bar"}]}))
+        Traceback (most recent call last):
+            ...
+        ValueError: Match revise config 0 must contain a _matcher key
+        >>> validate_match_revise(DictConfig({"_match_revise": [{"_matcher": {32: "bar"}}]}))
+        Traceback (most recent call last):
+            ...
+        ValueError: Match revise config 0 must contain a valid matcher in _matcher
+        >>> validate_match_revise(DictConfig({"_match_revise": [{"_matcher": {"code": "CODE//TEMP"}}]}))
+    """
+
+    if MATCH_REVISE_KEY not in stage_cfg:
+        raise ValueError(f"Stage configuration must contain a {MATCH_REVISE_KEY} key")
+
+    match_revise_options = stage_cfg[MATCH_REVISE_KEY]
+    if not isinstance(match_revise_options, (list, ListConfig)):
+        raise ValueError(f"Match revise options must be a list, got {type(match_revise_options)}")
+
+    for i, match_revise_cfg in enumerate(match_revise_options):
+        if not isinstance(match_revise_cfg, (dict, DictConfig)):
+            raise ValueError(f"Match revise config {i} must be a dict, got {type(match_revise_cfg)}")
+
+        if MATCHER_KEY not in match_revise_cfg:
+            raise ValueError(f"Match revise config {i} must contain a {MATCHER_KEY} key")
+
+        if not is_matcher(match_revise_cfg[MATCHER_KEY]):
+            raise ValueError(f"Match revise config {i} must contain a valid matcher in {MATCHER_KEY}")
+
+
+def match_revise_fntr(cfg: DictConfig, stage_cfg: DictConfig, compute_fn: ANY_COMPUTE_FN_T) -> COMPUTE_FN_T:
+    """A functor that creates a match & revise compute function based on the given configuration.
+
+    Stage configurations for match & revise must be in a match and revise format. Consider the below example,
+    showing the ``stage_cfg`` object in ``yaml`` format:
+
+        ```yaml
+        global_arg_1: "foo"
+        _match_revise:
+          - _matcher: {code: "CODE//BAR"}
+            local_arg_1: "bar"
+          - _matcher: {code: "CODE//BAZ"}
+            local_arg_1: "baz"
+        ```
+
+    This configuration will create a match & revise compute function that will filter the input DataFrame for
+    rows that match the ``CODE//BAR`` code and apply the compute function with the ``local_arg_1=bar``
+    parameter, and then filter the input DataFrame for rows that match the ``CODE//BAZ`` code and apply the
+    compute function with the ``local_arg_1=baz`` parameter. Both of these local compute functions will be
+    applied to the input DataFrame in sequence, and the resulting DataFrames will be concatenated alongside
+    any of the dataframe that matches no matcher (which will be left unmodified) and merged in a sorted way
+    that respects the ``patient_id``, ``time`` ordering first, then the order of the match & revise blocks
+    themselves, then the order of the rows in each match & revise block output. Each local compute function
+    will also use the ``global_arg_1=foo`` parameter.
+
+    Args:
+        cfg: The DictConfig configuration object.
+        stage_cfg: The DictConfig stage configuration object. This stage configuration must be in a match and
+            revise format, meaning it must have a key ``"_match_revise"`` that contains a list of local match
+            & revise configurations. Each local match & revise configuration must contain a key ``"_matcher"``
+            which links to the matcher configuration to use to filter the input DataFrame for the local
+            compute execution, and all other keys are local configuration parameters to be used in the local
+            compute execution.
+        compute_fn: The compute function to bind to the match & revise configuration local arguments.
+
+    Returns:
+        A function that applies the match & revise compute function to the input DataFrame.
+
+    Raises:
+        ValueError: If the stage configuration is not in a match and revise format.
+
+    Examples:
+        >>> df = pl.DataFrame({
+        ...     "patient_id": [1, 1, 1, 2, 2, 2],
+        ...     "time": [1, 2, 2, 1, 1, 2],
+        ...     "initial_idx": [0, 1, 2, 3, 4, 5],
+        ...     "code": ["FINAL", "CODE//TEMP_2", "CODE//TEMP_1", "FINAL", "CODE//TEMP_2", "CODE//TEMP_1"]
+        ... })
+        >>> def compute_fn(df: pl.DataFrame, stage_cfg: DictConfig) -> pl.DataFrame:
+        ...     return df.with_columns(
+        ...         pl.col("code").str.slice(0, len("CODE//")) +
+        ...         stage_cfg.local_code_mid + "//" + stage_cfg.global_code_end
+        ...     )
+        >>> stage_cfg = DictConfig({
+        ...     "global_code_end": "foo",
+        ...     "_match_revise": [
+        ...         {"_matcher": {"code": "CODE//TEMP_1"}, "local_code_mid": "bar"},
+        ...         {"_matcher": {"code": "CODE//TEMP_2"}, "local_code_mid": "baz"}
+        ...     ]
+        ... })
+        >>> cfg = DictConfig({"stage_cfg": stage_cfg})
+        >>> match_revise_fn = match_revise_fntr(cfg, stage_cfg, compute_fn)
+        >>> match_revise_fn(df.lazy()).collect()
+        shape: (6, 4)
+        ┌────────────┬──────┬─────────────┬────────────────┐
+        │ patient_id ┆ time ┆ initial_idx ┆ code           │
+        │ ---        ┆ ---  ┆ ---         ┆ ---            │
+        │ i64        ┆ i64  ┆ i64         ┆ str            │
+        ╞════════════╪══════╪═════════════╪════════════════╡
+        │ 1          ┆ 1    ┆ 0           ┆ FINAL          │
+        │ 1          ┆ 2    ┆ 2           ┆ CODE//bar//foo │
+        │ 1          ┆ 2    ┆ 1           ┆ CODE//baz//foo │
+        │ 2          ┆ 1    ┆ 4           ┆ CODE//baz//foo │
+        │ 2          ┆ 1    ┆ 3           ┆ FINAL          │
+        │ 2          ┆ 2    ┆ 5           ┆ CODE//bar//foo │
+        └────────────┴──────┴─────────────┴────────────────┘
+        >>> stage_cfg = DictConfig({
+        ...     "global_code_end": "foo",
+        ...     "_match_revise": [
+        ...         {"_matcher": {"code": "CODE//TEMP_2"}, "local_code_mid": "bizz"},
+        ...         {"_matcher": {"code": "CODE//TEMP_1"}, "local_code_mid": "foo", "global_code_end": "bar"},
+        ...     ]
+        ... })
+        >>> cfg = DictConfig({"stage_cfg": stage_cfg})
+        >>> match_revise_fn = match_revise_fntr(cfg, stage_cfg, compute_fn)
+        >>> match_revise_fn(df.lazy()).collect()
+        shape: (6, 4)
+        ┌────────────┬──────┬─────────────┬─────────────────┐
+        │ patient_id ┆ time ┆ initial_idx ┆ code            │
+        │ ---        ┆ ---  ┆ ---         ┆ ---             │
+        │ i64        ┆ i64  ┆ i64         ┆ str             │
+        ╞════════════╪══════╪═════════════╪═════════════════╡
+        │ 1          ┆ 1    ┆ 0           ┆ FINAL           │
+        │ 1          ┆ 2    ┆ 1           ┆ CODE//bizz//foo │
+        │ 1          ┆ 2    ┆ 2           ┆ CODE//foo//bar  │
+        │ 2          ┆ 1    ┆ 4           ┆ CODE//bizz//foo │
+        │ 2          ┆ 1    ┆ 3           ┆ FINAL           │
+        │ 2          ┆ 2    ┆ 5           ┆ CODE//foo//bar  │
+        └────────────┴──────┴─────────────┴─────────────────┘
+        >>> stage_cfg = DictConfig({
+        ...     "global_code_end": "foo", "_match_revise": [{"_matcher": {"missing": "CODE//TEMP_2"}}]
+        ... })
+        >>> cfg = DictConfig({"stage_cfg": stage_cfg})
+        >>> match_revise_fn = match_revise_fntr(cfg, stage_cfg, compute_fn)
+        >>> match_revise_fn(df.lazy()).collect() # doctest: +NORMALIZE_WHITESPACE
+        Traceback (most recent call last):
+            ...
+        ValueError: Missing needed columns {'missing'} for local matcher 0:
+            [(col("missing")) == (String(CODE//TEMP_2))].all_horizontal()
+        Columns available: 'code', 'initial_idx', 'patient_id', 'time'
+        >>> stage_cfg = DictConfig({"global_code_end": "foo"})
+        >>> cfg = DictConfig({"stage_cfg": stage_cfg})
+        >>> match_revise_fn = match_revise_fntr(cfg, stage_cfg, compute_fn)
+        Traceback (most recent call last):
+            ...
+        ValueError: Invalid match and revise configuration...
+    """
+    try:
+        validate_match_revise(stage_cfg)
+    except ValueError as e:
+        raise ValueError("Invalid match and revise configuration") from e
+
+    stage_cfg = hydra.utils.instantiate(stage_cfg)
+
+    matchers_and_fns = []
+    for match_revise_cfg in stage_cfg.pop(MATCH_REVISE_KEY):
+        matcher, cols = matcher_to_expr(match_revise_cfg.pop(MATCHER_KEY))
+        local_stage_cfg = DictConfig({**stage_cfg, **match_revise_cfg})
+        local_compute_fn = bind_compute_fn(cfg, local_stage_cfg, compute_fn)
+
+        matchers_and_fns.append((matcher, cols, local_compute_fn))
+
+    @wraps(compute_fn)
+    def match_revise_fn(df: DF_T) -> DF_T:
+        unmatched_df = df
+        cols = set(df.collect_schema().names())
+
+        revision_parts = []
+        for i, (matcher_expr, need_cols, local_compute_fn) in enumerate(matchers_and_fns):
+            if not need_cols.issubset(cols):
+                cols_str = "', '".join(x for x in sorted(cols))
+                raise ValueError(
+                    f"Missing needed columns {need_cols - cols} for local matcher {i}: "
+                    f"{matcher_expr}\nColumns available: '{cols_str}'"
+                )
+            matched_df = unmatched_df.filter(matcher_expr)
+            unmatched_df = unmatched_df.filter(~matcher_expr)
+
+            revision_parts.append(local_compute_fn(matched_df))
+
+        revision_parts.append(unmatched_df)
+        return pl.concat(revision_parts, how="vertical").sort(["patient_id", "time"], maintain_order=True)
+
+    return match_revise_fn
+
+
+def bind_compute_fn(cfg: DictConfig, stage_cfg: DictConfig, compute_fn: ANY_COMPUTE_FN_T) -> COMPUTE_FN_T:
+    """Bind the compute function to the appropriate parameters based on the type of the compute function.
+
+    Args:
+        cfg: The DictConfig configuration object.
+        stage_cfg: The DictConfig stage configuration object. This is separated from the ``cfg`` argument
+            because in some cases, such as under the match & revise paradigm, the stage config may be modified
+            dynamically under different matcher conditions to yield different compute functions.
+        compute_fn: The compute function to bind.
+
+    Returns:
+        The compute function bound to the appropriate parameters.
+
+    Raises:
+        ValueError: If the compute function is not a valid compute function.
+
+    Examples:
+        >>> compute_fn = bind_compute_fn(DictConfig({}), DictConfig({}), None)
+        >>> compute_fn("foobar")
+        'foobar'
+        >>> def compute_fntr(df: pl.DataFrame, cfg: DictConfig) -> pl.DataFrame:
+        ...     return df.with_columns(pl.lit(cfg.val).alias("val"))
+        >>> compute_fn = bind_compute_fn(DictConfig({"val": "foo"}), None, compute_fntr)
+        >>> compute_fn(pl.DataFrame({"a": [1, 2, 3]}))
+        shape: (3, 2)
+        ┌─────┬─────┐
+        │ a   ┆ val │
+        │ --- ┆ --- │
+        │ i64 ┆ str │
+        ╞═════╪═════╡
+        │ 1   ┆ foo │
+        │ 2   ┆ foo │
+        │ 3   ┆ foo │
+        └─────┴─────┘
+        >>> def compute_fntr(cfg: DictConfig) -> Callable[[pl.DataFrame], pl.DataFrame]:
+        ...     return lambda df: df.with_columns(pl.lit(cfg.val).alias("val"))
+        >>> compute_fn = bind_compute_fn(DictConfig({"val": "foo"}), None, compute_fntr)
+        >>> compute_fn(pl.DataFrame({"a": [1, 2, 3]}))
+        shape: (3, 2)
+        ┌─────┬─────┐
+        │ a   ┆ val │
+        │ --- ┆ --- │
+        │ i64 ┆ str │
+        ╞═════╪═════╡
+        │ 1   ┆ foo │
+        │ 2   ┆ foo │
+        │ 3   ┆ foo │
+        └─────┴─────┘
+        >>> def compute_fntr(stage_cfg, cfg) -> Callable[[pl.DataFrame], pl.DataFrame]:
+        ...     return lambda df: df.with_columns(
+        ...         pl.lit(stage_cfg.val).alias("stage_val"), pl.lit(cfg.val).alias("cfg_val")
+        ...     )
+        >>> compute_fn = bind_compute_fn(DictConfig({"val": "quo"}), DictConfig({"val": "bar"}), compute_fntr)
+        >>> compute_fn(pl.DataFrame({"a": [1, 2, 3]}))
+        shape: (3, 3)
+        ┌─────┬───────────┬─────────┐
+        │ a   ┆ stage_val ┆ cfg_val │
+        │ --- ┆ ---       ┆ ---     │
+        │ i64 ┆ str       ┆ str     │
+        ╞═════╪═══════════╪═════════╡
+        │ 1   ┆ bar       ┆ quo     │
+        │ 2   ┆ bar       ┆ quo     │
+        │ 3   ┆ bar       ┆ quo     │
+        └─────┴───────────┴─────────┘
+        >>> def compute_fntr(df, code_metadata):
+        ...     return df.join(code_metadata, on="a")
+        >>> code_metadata_df = pl.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
+        >>> from tempfile import TemporaryDirectory
+        >>> with TemporaryDirectory() as tmpdir:
+        ...     code_metadata_fp = Path(tmpdir) / "codes.parquet"
+        ...     code_metadata_df.write_parquet(code_metadata_fp)
+        ...     stage_cfg = DictConfig({"metadata_input_dir": tmpdir})
+        ...     compute_fn = bind_compute_fn(DictConfig({}), stage_cfg, compute_fntr)
+        ...     compute_fn(pl.DataFrame({"a": [1, 2, 3]}))
+        shape: (3, 2)
+        ┌─────┬─────┐
+        │ a   ┆ b   │
+        │ --- ┆ --- │
+        │ i64 ┆ i64 │
+        ╞═════╪═════╡
+        │ 1   ┆ 4   │
+        │ 2   ┆ 5   │
+        │ 3   ┆ 6   │
+        └─────┴─────┘
+        >>> def compute_fntr(df: pl.DataFrame, cfg: DictConfig) -> Callable[[pl.DataFrame], pl.DataFrame]:
+        ...     return lambda df: df
+        >>> bind_compute_fn(DictConfig({}), DictConfig({}), compute_fntr)
+        Traceback (most recent call last):
+            ...
+        ValueError: Invalid compute function
+    """
+
+    def fntr_params(compute_fn: ANY_COMPUTE_FN_T) -> ComputeFnArgs:
+        compute_fn_params = inspect.signature(compute_fn).parameters
+        kwargs = ComputeFnArgs()
+
+        if "cfg" in compute_fn_params:
+            kwargs["cfg"] = cfg
+        if "stage_cfg" in compute_fn_params:
+            kwargs["stage_cfg"] = stage_cfg
+        if "code_modifiers" in compute_fn_params:
+            code_modifiers = cfg.get("code_modifiers", None)
+            kwargs["code_modifiers"] = code_modifiers
+        if "code_metadata" in compute_fn_params:
+            kwargs["code_metadata"] = pl.read_parquet(
+                Path(stage_cfg.metadata_input_dir) / "codes.parquet", use_pyarrow=True
+            )
+        return kwargs
+
+    if compute_fn is None:
+        return identity_fn
+    match compute_fn_type(compute_fn):
+        case ComputeFnType.DIRECT:
+            pass
+        case ComputeFnType.UNBOUND:
+            compute_fn = partial(compute_fn, **fntr_params(compute_fn))
+        case ComputeFnType.FUNCTOR:
+            compute_fn = compute_fn(**fntr_params(compute_fn))
+        case _:
+            raise ValueError("Invalid compute function")
+
+    return compute_fn
+
+
 def map_over(
     cfg: DictConfig,
     compute_fn: COMPUTE_FN_T | None = None,
@@ -185,34 +607,10 @@ def map_over(
     elif includes_only_train:
         raise ValueError("All splits should be used, but shard iterator is returning only train splits?!?")
 
-    def fntr_params(compute_fn: ANY_COMPUTE_FN_T) -> ComputeFnArgs:
-        compute_fn_params = inspect.signature(compute_fn).parameters
-        kwargs = ComputeFnArgs()
-
-        if "cfg" in compute_fn_params:
-            kwargs["cfg"] = cfg
-        if "stage_cfg" in compute_fn_params:
-            kwargs["stage_cfg"] = cfg.stage_cfg
-        if "code_modifiers" in compute_fn_params:
-            code_modifiers = cfg.get("code_modifiers", None)
-            kwargs["code_modifiers"] = code_modifiers
-        if "code_metadata" in compute_fn_params:
-            kwargs["code_metadata"] = pl.read_parquet(
-                Path(cfg.stage_cfg.metadata_input_dir) / "codes.parquet", use_pyarrow=True
-            )
-        return kwargs
-
-    if compute_fn is None:
-        compute_fn = identity_fn
-    match compute_fn_type(compute_fn):
-        case ComputeFnType.DIRECT:
-            pass
-        case ComputeFnType.UNBOUND:
-            compute_fn = partial(compute_fn, **fntr_params(compute_fn))
-        case ComputeFnType.FUNCTOR:
-            compute_fn = compute_fn(**fntr_params(compute_fn))
-        case _:
-            raise ValueError("Invalid compute function")
+    if is_match_revise(cfg.stage_cfg):
+        compute_fn = match_revise_fntr(cfg, cfg.stage_cfg, compute_fn)
+    else:
+        compute_fn = bind_compute_fn(cfg, cfg.stage_cfg, compute_fn)
 
     all_out_fps = []
     for in_fp, out_fp in shards:


Test patch:
diff --git a/tests/test_add_time_derived_measurements.py b/tests/test_add_time_derived_measurements.py
index aad8428..87138c8 100644
--- a/tests/test_add_time_derived_measurements.py
+++ b/tests/test_add_time_derived_measurements.py
@@ -5,10 +5,7 @@ scripts.
 """
 
 
-from .transform_tester_base import (
-    ADD_TIME_DERIVED_MEASUREMENTS_SCRIPT,
-    single_stage_transform_tester,
-)
+from .transform_tester_base import ADD_TIME_DERIVED_MEASUREMENTS_SCRIPT, single_stage_transform_tester
 from .utils import parse_meds_csvs
 
 AGE_CALCULATION_STR = """
diff --git a/tests/test_filter_measurements.py b/tests/test_filter_measurements.py
index 8e0ff7e..2243320 100644
--- a/tests/test_filter_measurements.py
+++ b/tests/test_filter_measurements.py
@@ -5,10 +5,7 @@ scripts.
 """
 
 
-from .transform_tester_base import (
-    FILTER_MEASUREMENTS_SCRIPT,
-    single_stage_transform_tester,
-)
+from .transform_tester_base import FILTER_MEASUREMENTS_SCRIPT, single_stage_transform_tester
 from .utils import parse_meds_csvs
 
 # This is the code metadata
@@ -118,3 +115,114 @@ def test_filter_measurements():
         transform_stage_kwargs={"min_patients_per_code": 2},
         want_outputs=WANT_SHARDS,
     )
+
+
+# This is the code metadata
+# MEDS_CODE_METADATA_CSV = """
+# code,code/n_occurrences,code/n_patients,values/n_occurrences,values/sum,values/sum_sqd,description,parent_code
+# ,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
+# """
+#
+# In the test that applies to the match and revise framework, we'll filter codes in the following manner:
+#   - Codes that start with ADMISSION// will be filtered to occur at least 2 times, which are:
+#     ADMISSION//CARDIAC
+#   - Codes in [HR] will be filtered to occur at least 15 times, which are:
+#     (no codes)
+#   - Codes that start with EYE_COLOR// will be filtered to occur at least 4 times, which are:
+#     (no codes)
+#   - Other codes won't be filtered, so we will retain HEIGHT, DISCHARGE, DOB, TEMP
+
+MR_WANT_TRAIN_0 = """
+patient_id,time,code,numeric_value
+239684,,HEIGHT,175.271115221764
+239684,"12/28/1980, 00:00:00",DOB,
+239684,"05/11/2010, 17:41:51",ADMISSION//CARDIAC,
+239684,"05/11/2010, 17:41:51",TEMP,96.0
+239684,"05/11/2010, 17:48:48",TEMP,96.2
+239684,"05/11/2010, 18:25:35",TEMP,95.8
+239684,"05/11/2010, 18:57:18",TEMP,95.5
+239684,"05/11/2010, 19:27:19",DISCHARGE,
+1195293,,HEIGHT,164.6868838269085
+1195293,"06/20/1978, 00:00:00",DOB,
+1195293,"06/20/2010, 19:23:52",ADMISSION//CARDIAC,
+1195293,"06/20/2010, 19:23:52",TEMP,100.0
+1195293,"06/20/2010, 19:25:32",TEMP,100.0
+1195293,"06/20/2010, 19:45:19",TEMP,99.9
+1195293,"06/20/2010, 20:12:31",TEMP,99.8
+1195293,"06/20/2010, 20:24:44",TEMP,100.0
+1195293,"06/20/2010, 20:41:33",TEMP,100.4
+1195293,"06/20/2010, 20:50:04",DISCHARGE,
+"""
+
+MR_WANT_TRAIN_1 = """
+patient_id,time,code,numeric_value
+68729,,HEIGHT,160.3953106166676
+68729,"03/09/1978, 00:00:00",DOB,
+68729,"05/26/2010, 02:30:56",TEMP,97.8
+68729,"05/26/2010, 04:51:52",DISCHARGE,
+814703,,HEIGHT,156.48559093209357
+814703,"03/28/1976, 00:00:00",DOB,
+814703,"02/05/2010, 05:55:39",TEMP,100.1
+814703,"02/05/2010, 07:02:30",DISCHARGE,
+"""
+
+MR_WANT_TUNING_0 = """
+patient_id,time,code,numeric_value
+754281,,HEIGHT,166.22261567137025
+754281,"12/19/1988, 00:00:00",DOB,
+754281,"01/03/2010, 06:27:59",TEMP,99.8
+754281,"01/03/2010, 08:22:13",DISCHARGE,
+"""
+
+MR_WANT_HELD_OUT_0 = """
+patient_id,time,code,numeric_value
+1500733,,HEIGHT,158.60131573580904
+1500733,"07/20/1986, 00:00:00",DOB,
+1500733,"06/03/2010, 14:54:38",TEMP,100.0
+1500733,"06/03/2010, 15:39:49",TEMP,100.3
+1500733,"06/03/2010, 16:20:49",TEMP,100.1
+1500733,"06/03/2010, 16:44:26",DISCHARGE,
+"""
+
+MR_WANT_SHARDS = parse_meds_csvs(
+    {
+        "train/0": MR_WANT_TRAIN_0,
+        "train/1": MR_WANT_TRAIN_1,
+        "tuning/0": MR_WANT_TUNING_0,
+        "held_out/0": MR_WANT_HELD_OUT_0,
+    }
+)
+
+MATCH_REVISE_KEY = "_match_revise"
+MATCHER_KEY = "_matcher"
+
+
+def test_match_revise_filter_measurements():
+    single_stage_transform_tester(
+        transform_script=FILTER_MEASUREMENTS_SCRIPT,
+        stage_name="filter_measurements",
+        transform_stage_kwargs={
+            "_match_revise": [
+                {"_matcher": {"code": "ADMISSION//CARDIAC"}, "min_patients_per_code": 2},
+                {"_matcher": {"code": "ADMISSION//ORTHOPEDIC"}, "min_patients_per_code": 2},
+                {"_matcher": {"code": "ADMISSION//PULMONARY"}, "min_patients_per_code": 2},
+                {"_matcher": {"code": "HR"}, "min_patients_per_code": 15},
+                {"_matcher": {"code": "EYE_COLOR//BLUE"}, "min_patients_per_code": 4},
+                {"_matcher": {"code": "EYE_COLOR//BROWN"}, "min_patients_per_code": 4},
+                {"_matcher": {"code": "EYE_COLOR//HAZEL"}, "min_patients_per_code": 4},
+            ],
+        },
+        want_outputs=MR_WANT_SHARDS,
+        do_use_config_yaml=True,
+    )
diff --git a/tests/test_occlude_outliers.py b/tests/test_occlude_outliers.py
index 7eaeb26..2060bb7 100644
--- a/tests/test_occlude_outliers.py
+++ b/tests/test_occlude_outliers.py
@@ -7,10 +7,7 @@ scripts.
 
 import polars as pl
 
-from .transform_tester_base import (
-    OCCLUDE_OUTLIERS_SCRIPT,
-    single_stage_transform_tester,
-)
+from .transform_tester_base import OCCLUDE_OUTLIERS_SCRIPT, single_stage_transform_tester
 from .utils import MEDS_PL_SCHEMA, parse_meds_csvs
 
 # This is the code metadata
diff --git a/tests/test_reorder_measurements.py b/tests/test_reorder_measurements.py
index 2f3dd0b..305b85b 100644
--- a/tests/test_reorder_measurements.py
+++ b/tests/test_reorder_measurements.py
@@ -5,10 +5,7 @@ scripts.
 """
 
 
-from .transform_tester_base import (
-    REORDER_MEASUREMENTS_SCRIPT,
-    single_stage_transform_tester,
-)
+from .transform_tester_base import REORDER_MEASUREMENTS_SCRIPT, single_stage_transform_tester
 from .utils import parse_meds_csvs
 
 ORDERED_CODE_PATTERNS = [
diff --git a/tests/transform_tester_base.py b/tests/transform_tester_base.py
index ce2b8c3..80e4f4b 100644
--- a/tests/transform_tester_base.py
+++ b/tests/transform_tester_base.py
@@ -291,6 +291,7 @@ def single_stage_transform_tester(
     input_shards: dict[str, pl.DataFrame] | None = None,
     do_pass_stage_name: bool = False,
     file_suffix: str = ".parquet",
+    do_use_config_yaml: bool = False,
 ):
     with tempfile.TemporaryDirectory() as d:
         MEDS_dir = Path(d) / "MEDS_cohort"
@@ -337,12 +338,17 @@ def single_stage_transform_tester(
         if transform_stage_kwargs:
             pipeline_config_kwargs["stage_configs"] = {stage_name: transform_stage_kwargs}
 
+        run_command_kwargs = {
+            "script": transform_script,
+            "hydra_kwargs": pipeline_config_kwargs,
+            "test_name": f"Single stage transform: {stage_name}",
+        }
+        if do_use_config_yaml:
+            run_command_kwargs["do_use_config_yaml"] = True
+            run_command_kwargs["config_name"] = "preprocess"
+
         # Run the transform
-        stderr, stdout = run_command(
-            transform_script,
-            pipeline_config_kwargs,
-            f"Single stage transform: {stage_name}",
-        )
+        stderr, stdout = run_command(**run_command_kwargs)
 
         # Check the output
         if isinstance(want_outputs, pl.DataFrame):
diff --git a/tests/utils.py b/tests/utils.py
index 2f383f0..d6ac438 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -1,8 +1,10 @@
 import subprocess
+import tempfile
 from io import StringIO
 from pathlib import Path
 
 import polars as pl
+from omegaconf import OmegaConf
 from polars.testing import assert_frame_equal
 
 DEFAULT_CSV_TS_FORMAT = "%m/%d/%Y, %H:%M:%S"
@@ -102,31 +104,66 @@ def run_command(
     test_name: str,
     config_name: str | None = None,
     should_error: bool = False,
+    do_use_config_yaml: bool = False,
 ):
     script = ["python", str(script.resolve())] if isinstance(script, Path) else [script]
     command_parts = script
-    if config_name is not None:
-        command_parts.append(f"--config-name={config_name}")
-    command_parts.append(" ".join(dict_to_hydra_kwargs(hydra_kwargs)))
+
+    err_cmd_lines = []
+
+    if do_use_config_yaml:
+        if config_name is None:
+            raise ValueError("config_name must be provided if do_use_config_yaml is True.")
+
+        conf = OmegaConf.create(
+            {
+                "defaults": [config_name],
+                **hydra_kwargs,
+            }
+        )
+
+        conf_dir = tempfile.TemporaryDirectory()
+        conf_path = Path(conf_dir.name) / "config.yaml"
+        OmegaConf.save(conf, conf_path)
+
+        command_parts.extend(
+            [
+                f"--config-path={str(conf_path.parent.resolve())}",
+                "--config-name=config",
+                "'hydra.searchpath=[pkg://MEDS_transforms.configs]'",
+            ]
+        )
+        err_cmd_lines.append(f"Using config yaml:\n{OmegaConf.to_yaml(conf)}")
+    else:
+        if config_name is not None:
+            command_parts.append(f"--config-name={config_name}")
+        command_parts.append(" ".join(dict_to_hydra_kwargs(hydra_kwargs)))
 
     full_cmd = " ".join(command_parts)
+    err_cmd_lines.append(f"Running command: {full_cmd}")
     command_out = subprocess.run(full_cmd, shell=True, capture_output=True)
 
     command_errored = command_out.returncode != 0
 
     stderr = command_out.stderr.decode()
+    err_cmd_lines.append(f"stderr:\n{stderr}")
     stdout = command_out.stdout.decode()
+    err_cmd_lines.append(f"stdout:\n{stdout}")
 
     if should_error and not command_errored:
+        if do_use_config_yaml:
+            conf_dir.cleanup()
         raise AssertionError(
-            f"{test_name} failed as command did not error when expected!\n"
-            f"command:{full_cmd}\nstdout:\n{stdout}\nstderr:\n{stderr}"
+            f"{test_name} failed as command did not error when expected!\n" + "\n".join(err_cmd_lines)
         )
     elif not should_error and command_errored:
+        if do_use_config_yaml:
+            conf_dir.cleanup()
         raise AssertionError(
-            f"{test_name} failed as command errored when not expected!"
-            f"\ncommand:{full_cmd}\nstdout:\n{stdout}\nstderr:\n{stderr}"
+            f"{test_name} failed as command errored when not expected!\n" + "\n".join(err_cmd_lines)
         )
+    if do_use_config_yaml:
+        conf_dir.cleanup()
     return stderr, stdout