Collator

The Collator is the first stage of the pipeline. It collects the many raw tables that make up a dataset (parquet or csv) and collates them into a single, denormalized event stream in a MEDS-like format: one row per event, with a uniform schema of subject_id, time, code, numeric_value, and text_value.

Which tables to read, how to filter and aggregate them, and how each event maps onto a code are all driven by a collation config (collation.yaml). Timestamps are normalized to timezone-naive UTC along the way.

What it produces

Running the collator (via save_all) writes two files to the processed-data directory:

• meds.parquet — the collated events for every subject.
• subject_splits.parquet — a train / tuning / held_out assignment for each subject, partitioned by first event time and the fractions in the config.

How it works

1. get_reference_frame builds a per-subject "spine" (for example, a hospital stay with a start and end time), optionally joining in augmentation tables.
2. get_entry turns one configured event type into standardized token rows — building the code string, casting values, and optionally restricting events to the reference window.
3. get_all concatenates every configured entry into the full event stream.
4. get_subject_splits assigns the data splits.

!!! warning "Config files are trusted input" Collation configs may contain Polars expression strings (for filters, added columns, and aggregations) that are evaluated at load time. These run through slightly_safer_eval, which restricts the available namespace but is not secure against malicious input. Only run configs you trust.

collects and collates different dataframes into a denormalized format

Collator

Bases: Configurable

Source code in src/cocoa/collator.py

class Collator(Configurable):
    default_file = "collation.yaml"

    def __init__(
        self,
        collation_cfg: pathlib.Path | str = None,
        raw_data_home: pathlib.Path | str = None,
        processed_data_home: pathlib.Path | str = None,
        **kwargs,
    ):
        super().__init__(collation_cfg, **kwargs)

        self.raw_data_home, self.processed_data_home = map(
            lambda p: pathlib.Path(p).expanduser().resolve(),
            [raw_data_home, processed_data_home],
        )
        self.processed_data_home.mkdir(parents=True, exist_ok=True)
        self.reference_frame = None
        self.splits: tuple = ("train", "tuning", "held_out")
        self._tz_warned: set = set()

        self.logger.info("Collator initialized...")
        self.logger.info(f"{self.raw_data_home=}")
        self.logger.info(f"{self.processed_data_home=}")

    @staticmethod
    def slightly_safer_eval(expr):
        """
        performs eval on polars expressions;
        prevents some forms of accidental usage
        but is not secure against malicious input
        """
        return eval(
            expr,
            {"__builtins__": {"str": str, "int": int, "float": float, "bool": bool}},
            {"pl": pl},
        )

    def to_utc_naive(
        self, df: pl.LazyFrame, column: str, *, table: str = None
    ) -> pl.Expr:
        """
        expression converting `column` to a timezone-naive UTC datetime;
        tz-aware columns are converted to UTC (instant-preserving),
        tz-naive columns are assumed to already be UTC
        """
        dtype = df.collect_schema().get(column)
        if isinstance(dtype, pl.Datetime) and dtype.time_zone not in (None, "UTC"):
            if (key := (table, column, dtype.time_zone)) not in self._tz_warned:
                self._tz_warned.add(key)
                self.logger.warning(
                    f"{table or '<frame>'}.{column}: "
                    f"converting {dtype.time_zone} -> UTC"
                )
            return (
                pl.col(column)
                .dt.convert_time_zone("UTC")
                .dt.replace_time_zone(time_zone=None)
            )
        return pl.col(column).cast(pl.Datetime).dt.replace_time_zone(time_zone=None)

    def load_table(
        self,
        *,
        table: str = None,
        filter_expr: str | list = None,
        agg_expr: str | list = None,
        with_col_expr: str | list = None,
        key: str = None,
        subject_id_str: str = None,
        time: str = None,
        **kwargs,
    ) -> pl.LazyFrame:
        """
        lazy-load the table `table`.parquet and perform some standard ETL
        tasks in the following order if specified:
        1. fix subject_id to `subject_id_str`
        2. perform a filter operation
        3. add columns
        4. perform an aggregation by `key` or self.cfg["subfject_id"]
        """
        if (f := self.raw_data_home / f"{table}.parquet").exists():
            df = pl.scan_parquet(f)
        elif (f := self.raw_data_home / f"{table}.csv").exists():
            df = pl.scan_csv(f)
        else:
            raise FileNotFoundError(
                f"No parquet / csv file found for {table=} in {self.raw_data_home}"
            )
        if subject_id_str is not None:
            df = df.with_columns(pl.col(subject_id_str).alias(self.cfg["subject_id"]))
        if filter_expr is not None:
            df = df.filter(
                self.slightly_safer_eval(filter_expr)
                if isinstance(filter_expr, str)
                else [self.slightly_safer_eval(c) for c in filter_expr]
            )
        if with_col_expr is not None:
            df = df.with_columns(
                self.slightly_safer_eval(with_col_expr)
                if isinstance(with_col_expr, str)
                else [self.slightly_safer_eval(c) for c in with_col_expr]
            )
        if agg_expr is not None:
            df = (
                df.sort(time)
                .group_by(key if key is not None else self.cfg["subject_id"])
                .agg(
                    self.slightly_safer_eval(agg_expr)
                    if isinstance(agg_expr, str)
                    else [self.slightly_safer_eval(c) for c in agg_expr]
                )
            )
        return df

    def get_reference_frame(self) -> pl.LazyFrame:
        """create the static reference frame as configured"""
        cfg = self.cfg["reference"]
        if self.reference_frame is not None:  # pull from cache if available
            return self.reference_frame
        df = self.load_table(**cfg)
        for tkv in cfg.get("augmentation_tables", []):
            df = df.join(
                self.load_table(**tkv),
                on=tkv["key"],
                validate=tkv["validation"],
                how="left",
                maintain_order="left",
            )
        for col in (cfg["start_time"], cfg["end_time"]):
            df = df.with_columns(
                self.to_utc_naive(df, col, table=cfg["table"]).alias(col)
            )
        self.reference_frame = df  # cache result
        return self.reference_frame

    def get_entry(
        self,
        table: str,
        time: str,
        code: str,
        *,
        prefix: str = None,
        numeric_value: str = None,
        text_value: str = None,
        filter_expr: str = None,
        with_col_expr: str = None,
        reference_key: str = None,
        subject_id_str: str = None,
        fix_date_to_time: bool = None,
        agg_expr: str | list = None,
        key: str = None,
    ) -> pl.LazyFrame:
        """create tokens corresponding to a configured event"""
        df = (
            self.get_reference_frame()
            if table == "REFERENCE"
            else self.load_table(
                table=table,
                time=time,
                filter_expr=filter_expr,
                with_col_expr=with_col_expr,
                subject_id_str=subject_id_str,
                agg_expr=agg_expr,
                key=key,
            )
        )
        if time in df.collect_schema().names():
            df = df.with_columns(self.to_utc_naive(df, time, table=table).alias(time))
        # otherwise `time` arrives normalized via the reference_key join below
        if fix_date_to_time:
            # if a date was cast to a time,
            # the default of 00:00:00 should be replaced with 23:59:59
            df = df.with_columns(
                pl.when(pl.col(time).cast(pl.Datetime).dt.time() == pl.time(0, 0, 0))
                .then(
                    pl.col(time)
                    .cast(pl.Datetime)
                    .dt.replace(hour=23, minute=59, second=59)
                )
                .otherwise(pl.col(time).cast(pl.Datetime))
                .alias(time)
            )
        if reference_key is not None:
            df = df.join(self.reference_frame, on=reference_key, how="inner").filter(
                pl.col(time)
                .cast(pl.Datetime)
                .is_between(
                    pl.col(self.cfg["reference"]["start_time"]).cast(pl.Datetime),
                    pl.col(self.cfg["reference"]["end_time"]).cast(pl.Datetime),
                )
            )

        return df.select(
            pl.col(self.cfg["subject_id"]).cast(pl.String).alias("subject_id"),
            # `time` was normalized to naive UTC when the frame was loaded above
            pl.col(time).alias("time"),
            pl.when(pl.col(code).is_not_null())
            .then(
                pl.concat_str(
                    [
                        pl.lit(prefix),
                        pl.col(code)
                        .cast(pl.String)
                        .str.to_lowercase()
                        .str.replace_all(r"\s+", "_"),
                    ],
                    separator="//",
                    ignore_nulls=True,  # prefix is optional
                )
            )
            .alias("code"),
            (pl.col(numeric_value) if numeric_value else pl.lit(None))
            .cast(pl.Float32)  # dumb
            .alias("numeric_value"),
            (
                pl.col(text_value).str.to_lowercase().str.replace_all(r"\s+", "_")
                if text_value
                else pl.lit(None)
            )
            .cast(pl.String)
            .alias("text_value"),
        ).drop_nulls(subset=["subject_id", "time", "code"])

    def get_all(self) -> pl.LazyFrame:
        """get all tokens for all events as configured"""
        return pl.concat(
            (self.get_entry(**entry) for entry in self.cfg.get("entries", []))
        )

    def get_subject_splits(self) -> pl.DataFrame:
        """get the subject splits as configured"""
        partition = (
            self.get_reference_frame()
            .group_by(self.cfg.get("group_id", self.cfg["subject_id"]))
            .agg(pl.col(self.cfg.reference.start_time).min().alias("first_time"))
            .sort("first_time")
            .with_row_index()
        ).collect()
        split_idx = (
            (
                len(partition)
                * np.cumsum(
                    [
                        self.cfg.subject_splits.train_frac,
                        self.cfg.subject_splits.tuning_frac,
                    ]
                )
            )
            .astype(int)
            .tolist()
        )
        partition = partition.with_columns(
            split=pl.when(pl.col("index") < split_idx[0])
            .then(pl.lit(self.splits[0]))
            .when(pl.col("index") < split_idx[1])
            .then(pl.lit(self.splits[1]))
            .otherwise(pl.lit(self.splits[2]))
        )
        return (
            partition
            if "group_id" not in self.cfg
            else self.get_reference_frame()
            .collect()
            .join(partition, on=self.cfg["group_id"])
        ).select(pl.col(self.cfg["subject_id"]).alias("subject_id"), "split")

    def save_all(self, verbose: bool = False):
        """save collated data and subject splits to disc, optionally w/ summary stats"""
        meds_path = self.processed_data_home / "meds.parquet"
        meds_path.unlink(missing_ok=True)
        df_all = self.get_all()
        try:
            df_all.sink_parquet(meds_path, engine="streaming")
        except Exception as e:
            self.logger.warning(f"Streaming write failed: {e}")
            df_all.sink_parquet(meds_path, engine="in-memory", mkdir=True)
        df_splits = self.get_subject_splits()
        if "pass_through_columns" in self.cfg:
            df_splits = df_splits.join(
                self.reference_frame.select(
                    pl.col(self.cfg["subject_id"]).alias("subject_id"),
                    *self.cfg.pass_through_columns,
                ).collect(),
                on="subject_id",
                how="inner",
                validate="1:1",
            )
        df_splits.write_parquet(
            self.processed_data_home / "subject_splits.parquet", mkdir=True
        )

        if verbose:
            self.logger.summarize_meds_like(df_all, df_splits)

get_all()

get all tokens for all events as configured

Source code in src/cocoa/collator.py

def get_all(self) -> pl.LazyFrame:
    """get all tokens for all events as configured"""
    return pl.concat(
        (self.get_entry(**entry) for entry in self.cfg.get("entries", []))
    )

get_entry(table, time, code, *, prefix=None, numeric_value=None, text_value=None, filter_expr=None, with_col_expr=None, reference_key=None, subject_id_str=None, fix_date_to_time=None, agg_expr=None, key=None)

create tokens corresponding to a configured event

Source code in src/cocoa/collator.py

def get_entry(
    self,
    table: str,
    time: str,
    code: str,
    *,
    prefix: str = None,
    numeric_value: str = None,
    text_value: str = None,
    filter_expr: str = None,
    with_col_expr: str = None,
    reference_key: str = None,
    subject_id_str: str = None,
    fix_date_to_time: bool = None,
    agg_expr: str | list = None,
    key: str = None,
) -> pl.LazyFrame:
    """create tokens corresponding to a configured event"""
    df = (
        self.get_reference_frame()
        if table == "REFERENCE"
        else self.load_table(
            table=table,
            time=time,
            filter_expr=filter_expr,
            with_col_expr=with_col_expr,
            subject_id_str=subject_id_str,
            agg_expr=agg_expr,
            key=key,
        )
    )
    if time in df.collect_schema().names():
        df = df.with_columns(self.to_utc_naive(df, time, table=table).alias(time))
    # otherwise `time` arrives normalized via the reference_key join below
    if fix_date_to_time:
        # if a date was cast to a time,
        # the default of 00:00:00 should be replaced with 23:59:59
        df = df.with_columns(
            pl.when(pl.col(time).cast(pl.Datetime).dt.time() == pl.time(0, 0, 0))
            .then(
                pl.col(time)
                .cast(pl.Datetime)
                .dt.replace(hour=23, minute=59, second=59)
            )
            .otherwise(pl.col(time).cast(pl.Datetime))
            .alias(time)
        )
    if reference_key is not None:
        df = df.join(self.reference_frame, on=reference_key, how="inner").filter(
            pl.col(time)
            .cast(pl.Datetime)
            .is_between(
                pl.col(self.cfg["reference"]["start_time"]).cast(pl.Datetime),
                pl.col(self.cfg["reference"]["end_time"]).cast(pl.Datetime),
            )
        )

    return df.select(
        pl.col(self.cfg["subject_id"]).cast(pl.String).alias("subject_id"),
        # `time` was normalized to naive UTC when the frame was loaded above
        pl.col(time).alias("time"),
        pl.when(pl.col(code).is_not_null())
        .then(
            pl.concat_str(
                [
                    pl.lit(prefix),
                    pl.col(code)
                    .cast(pl.String)
                    .str.to_lowercase()
                    .str.replace_all(r"\s+", "_"),
                ],
                separator="//",
                ignore_nulls=True,  # prefix is optional
            )
        )
        .alias("code"),
        (pl.col(numeric_value) if numeric_value else pl.lit(None))
        .cast(pl.Float32)  # dumb
        .alias("numeric_value"),
        (
            pl.col(text_value).str.to_lowercase().str.replace_all(r"\s+", "_")
            if text_value
            else pl.lit(None)
        )
        .cast(pl.String)
        .alias("text_value"),
    ).drop_nulls(subset=["subject_id", "time", "code"])

get_reference_frame()

create the static reference frame as configured

Source code in src/cocoa/collator.py

def get_reference_frame(self) -> pl.LazyFrame:
    """create the static reference frame as configured"""
    cfg = self.cfg["reference"]
    if self.reference_frame is not None:  # pull from cache if available
        return self.reference_frame
    df = self.load_table(**cfg)
    for tkv in cfg.get("augmentation_tables", []):
        df = df.join(
            self.load_table(**tkv),
            on=tkv["key"],
            validate=tkv["validation"],
            how="left",
            maintain_order="left",
        )
    for col in (cfg["start_time"], cfg["end_time"]):
        df = df.with_columns(
            self.to_utc_naive(df, col, table=cfg["table"]).alias(col)
        )
    self.reference_frame = df  # cache result
    return self.reference_frame

get_subject_splits()

get the subject splits as configured

Source code in src/cocoa/collator.py

def get_subject_splits(self) -> pl.DataFrame:
    """get the subject splits as configured"""
    partition = (
        self.get_reference_frame()
        .group_by(self.cfg.get("group_id", self.cfg["subject_id"]))
        .agg(pl.col(self.cfg.reference.start_time).min().alias("first_time"))
        .sort("first_time")
        .with_row_index()
    ).collect()
    split_idx = (
        (
            len(partition)
            * np.cumsum(
                [
                    self.cfg.subject_splits.train_frac,
                    self.cfg.subject_splits.tuning_frac,
                ]
            )
        )
        .astype(int)
        .tolist()
    )
    partition = partition.with_columns(
        split=pl.when(pl.col("index") < split_idx[0])
        .then(pl.lit(self.splits[0]))
        .when(pl.col("index") < split_idx[1])
        .then(pl.lit(self.splits[1]))
        .otherwise(pl.lit(self.splits[2]))
    )
    return (
        partition
        if "group_id" not in self.cfg
        else self.get_reference_frame()
        .collect()
        .join(partition, on=self.cfg["group_id"])
    ).select(pl.col(self.cfg["subject_id"]).alias("subject_id"), "split")

load_table(*, table=None, filter_expr=None, agg_expr=None, with_col_expr=None, key=None, subject_id_str=None, time=None, **kwargs)

lazy-load the table table.parquet and perform some standard ETL tasks in the following order if specified: 1. fix subject_id to subject_id_str 2. perform a filter operation 3. add columns 4. perform an aggregation by key or self.cfg["subfject_id"]

Source code in src/cocoa/collator.py

def load_table(
    self,
    *,
    table: str = None,
    filter_expr: str | list = None,
    agg_expr: str | list = None,
    with_col_expr: str | list = None,
    key: str = None,
    subject_id_str: str = None,
    time: str = None,
    **kwargs,
) -> pl.LazyFrame:
    """
    lazy-load the table `table`.parquet and perform some standard ETL
    tasks in the following order if specified:
    1. fix subject_id to `subject_id_str`
    2. perform a filter operation
    3. add columns
    4. perform an aggregation by `key` or self.cfg["subfject_id"]
    """
    if (f := self.raw_data_home / f"{table}.parquet").exists():
        df = pl.scan_parquet(f)
    elif (f := self.raw_data_home / f"{table}.csv").exists():
        df = pl.scan_csv(f)
    else:
        raise FileNotFoundError(
            f"No parquet / csv file found for {table=} in {self.raw_data_home}"
        )
    if subject_id_str is not None:
        df = df.with_columns(pl.col(subject_id_str).alias(self.cfg["subject_id"]))
    if filter_expr is not None:
        df = df.filter(
            self.slightly_safer_eval(filter_expr)
            if isinstance(filter_expr, str)
            else [self.slightly_safer_eval(c) for c in filter_expr]
        )
    if with_col_expr is not None:
        df = df.with_columns(
            self.slightly_safer_eval(with_col_expr)
            if isinstance(with_col_expr, str)
            else [self.slightly_safer_eval(c) for c in with_col_expr]
        )
    if agg_expr is not None:
        df = (
            df.sort(time)
            .group_by(key if key is not None else self.cfg["subject_id"])
            .agg(
                self.slightly_safer_eval(agg_expr)
                if isinstance(agg_expr, str)
                else [self.slightly_safer_eval(c) for c in agg_expr]
            )
        )
    return df

save_all(verbose=False)

save collated data and subject splits to disc, optionally w/ summary stats

Source code in src/cocoa/collator.py

def save_all(self, verbose: bool = False):
    """save collated data and subject splits to disc, optionally w/ summary stats"""
    meds_path = self.processed_data_home / "meds.parquet"
    meds_path.unlink(missing_ok=True)
    df_all = self.get_all()
    try:
        df_all.sink_parquet(meds_path, engine="streaming")
    except Exception as e:
        self.logger.warning(f"Streaming write failed: {e}")
        df_all.sink_parquet(meds_path, engine="in-memory", mkdir=True)
    df_splits = self.get_subject_splits()
    if "pass_through_columns" in self.cfg:
        df_splits = df_splits.join(
            self.reference_frame.select(
                pl.col(self.cfg["subject_id"]).alias("subject_id"),
                *self.cfg.pass_through_columns,
            ).collect(),
            on="subject_id",
            how="inner",
            validate="1:1",
        )
    df_splits.write_parquet(
        self.processed_data_home / "subject_splits.parquet", mkdir=True
    )

    if verbose:
        self.logger.summarize_meds_like(df_all, df_splits)

slightly_safer_eval(expr) staticmethod

performs eval on polars expressions; prevents some forms of accidental usage but is not secure against malicious input

Source code in src/cocoa/collator.py

@staticmethod
def slightly_safer_eval(expr):
    """
    performs eval on polars expressions;
    prevents some forms of accidental usage
    but is not secure against malicious input
    """
    return eval(
        expr,
        {"__builtins__": {"str": str, "int": int, "float": float, "bool": bool}},
        {"pl": pl},
    )

to_utc_naive(df, column, *, table=None)

expression converting column to a timezone-naive UTC datetime; tz-aware columns are converted to UTC (instant-preserving), tz-naive columns are assumed to already be UTC

Source code in src/cocoa/collator.py

def to_utc_naive(
    self, df: pl.LazyFrame, column: str, *, table: str = None
) -> pl.Expr:
    """
    expression converting `column` to a timezone-naive UTC datetime;
    tz-aware columns are converted to UTC (instant-preserving),
    tz-naive columns are assumed to already be UTC
    """
    dtype = df.collect_schema().get(column)
    if isinstance(dtype, pl.Datetime) and dtype.time_zone not in (None, "UTC"):
        if (key := (table, column, dtype.time_zone)) not in self._tz_warned:
            self._tz_warned.add(key)
            self.logger.warning(
                f"{table or '<frame>'}.{column}: "
                f"converting {dtype.time_zone} -> UTC"
            )
        return (
            pl.col(column)
            .dt.convert_time_zone("UTC")
            .dt.replace_time_zone(time_zone=None)
        )
    return pl.col(column).cast(pl.Datetime).dt.replace_time_zone(time_zone=None)