convert_to_subject_sharded

Stage: convert event-sharded raw data into subject-sharded format.

This is the second half of the initial ingestion phase. After convert_to_parquet has normalized each raw source table to parquet, this stage re-groups rows by subject: for every (split, table) pair, it reads every file of that table, applies the table’s subject_id expression (and any join it needs), filters down to the rows whose subject is in the split, and writes the result to <split>/<table>.parquet.

For example, with a vitals table joined to stays on stay_id:

.. code-block:: text

data/vitals/[0-2).parquet  ─┐
data/vitals/[2-4).parquet  ─┼─► data/train/0/vitals.parquet
data/vitals/[4-6).parquet  ─┘    (vitals rows for training subjects, with
                                  the joined subject_id materialized)

Each shard is independent, so this stage parallelizes trivially across (split, table) pairs.

_filter_to_subjects(df, *, table, subjects)

Filter df to the rows whose subject_id is in subjects.

Uses the table’s subject_id_polars_expr inline in filter so that the output keeps its original source columns untouched — no new materialized subject_id column gets added.

Source code in MEDS_extract/convert_to_subject_sharded/convert_to_subject_sharded.py

def _filter_to_subjects(df: pl.LazyFrame, *, table: TableConfig, subjects: Sequence[int]) -> pl.LazyFrame:
    """Filter ``df`` to the rows whose subject_id is in ``subjects``.

    Uses the table's ``subject_id_polars_expr`` inline in ``filter`` so that
    the output keeps its original source columns untouched — no new
    materialized ``subject_id`` column gets added.
    """
    return df.filter(table.subject_id_polars_expr.is_in(list(subjects)))

_read_and_join(fps, *, table, input_dir)

Scan the given subshards and apply the table’s join (if any).

This is a straight read — no filtering — so a row-level rwlock wrapper can reuse the same function across multiple stages. Subject filtering lives in :func:_filter_to_subjects on the compute side.

Source code in MEDS_extract/convert_to_subject_sharded/convert_to_subject_sharded.py

def _read_and_join(
    fps: Sequence[Path],
    *,
    table: TableConfig,
    input_dir: Path,
) -> pl.LazyFrame:
    """Scan the given subshards and apply the table's join (if any).

    This is a straight read — no filtering — so a row-level rwlock
    wrapper can reuse the same function across multiple stages. Subject
    filtering lives in :func:`_filter_to_subjects` on the compute side.
    """
    df = scan_source(fps)
    return table.apply_join(df, input_dir)

main(cfg)

Re-shard raw data by subject. See module docstring for details.

All arguments come through the Hydra cfg object; this stage has no stage-specific options beyond the global MESSY_config_fp.

Source code in MEDS_extract/convert_to_subject_sharded/convert_to_subject_sharded.py

@Stage.register(is_metadata=False, example_class=MEDSExtractStageExample)
def main(cfg: DictConfig):
    """Re-shard raw data by subject. See module docstring for details.

    All arguments come through the Hydra ``cfg`` object; this stage has no
    stage-specific options beyond the global ``MESSY_config_fp``.
    """
    input_dir = Path(cfg.stage_cfg.data_input_dir)
    subject_subsharded_dir = Path(cfg.stage_cfg.output_dir)

    shards = json.loads(Path(cfg.shards_map_fp).read_text())

    messy_cfg = MessyConfig.load(cfg.MESSY_config_fp)
    subject_subsharded_dir.mkdir(parents=True, exist_ok=True)

    subject_splits = list(shards.items())
    random.shuffle(subject_splits)

    for sp, subjects in subject_splits:
        for table in messy_cfg.shuffled_tables():
            event_shards = list(table.source_files(input_dir))
            random.shuffle(event_shards)

            out_fp = subject_subsharded_dir / sp / f"{table.input_prefix}.parquet"

            rwlock_wrap(
                event_shards,
                out_fp,
                partial(_read_and_join, table=table, input_dir=input_dir),
                sink_df,
                partial(_filter_to_subjects, table=table, subjects=subjects),
                do_overwrite=cfg.do_overwrite,
                # Run-scoped do_overwrite (MT 0.7.0): without the marker dir, parallel
                # workers treat each other's fresh outputs as stale and redo the work.
                marker_dir=run_marker_dir(cfg),
            )

    logger.info("Created a subject-sharded view.")

sink_df(df, out_fp)

Atomically sink a lazy plan to parquet on polars’ streaming engine.

The drop-in replacement for MEDS-transforms’ eager write_df on this stage — same .tmp + os.replace atomicity — but the plan executes streamed, so scan → join → subject-filter pipelines in chunks and peak memory is O(rows written to this shard), not O(full joined table) — measured 5.2x lower on a 30M-row table, with the eager path’s exact row order. Order determinism is load-bearing and doubly pinned: the ordered join in JoinConfig.apply (maintain_order="left_right") plus maintain_order=True here — without both, streaming execution reorders nondeterministically and merge’s stable sort would propagate that order into the final MEDS bytes for same-time events.

Source code in MEDS_extract/convert_to_subject_sharded/convert_to_subject_sharded.py

def sink_df(df: pl.LazyFrame, out_fp: Path) -> None:
    """Atomically sink a lazy plan to parquet on polars' streaming engine.

    The drop-in replacement for MEDS-transforms' eager ``write_df`` on this stage —
    same ``.tmp`` + ``os.replace`` atomicity — but the plan executes streamed, so
    scan → join → subject-filter pipelines in chunks and peak memory is O(rows
    written to this shard), not O(full joined table) — measured 5.2x lower on a
    30M-row table, with the eager path's exact row order. Order determinism is
    load-bearing and doubly pinned: the ordered join in ``JoinConfig.apply``
    (``maintain_order="left_right"``) plus ``maintain_order=True`` here — without
    both, streaming execution reorders nondeterministically and merge's stable sort
    would propagate that order into the final MEDS bytes for same-time events.
    """
    out_fp.parent.mkdir(parents=True, exist_ok=True)
    tmp_fp = out_fp.with_suffix(out_fp.suffix + ".tmp")
    try:
        df.lazy().sink_parquet(tmp_fp, maintain_order=True)
        os.replace(tmp_fp, out_fp)
    except BaseException:
        tmp_fp.unlink(missing_ok=True)
        raise