Making an analysis model fit in memory

A stretch of work on the models behind a bank's internal data analysis. The numbers were not in question. The problem was that a full run took long enough, and asked for enough memory, that people had quietly stopped re-running it and started working from extracts that were a few days old.

Where the time and memory went

Very little of the cost was arithmetic. Most of it was intermediate copies: a chain of transformations that each materialised a new frame, so peak memory tracked the length of the chain rather than the size of the data. Loading with narrower dtypes, dropping columns before the joins instead of after, and replacing the row-wise applies with vectorised expressions took care of most of it. Moving the heavier steps to a lazy frame meant the projection and filter pushdown happened without anyone having to hand-order the pipeline.

def daily_exposure(trades: pl.LazyFrame, fx: pl.LazyFrame) -> pl.LazyFrame:
    return (
        trades
        .select("trade_date", "book", "ccy", "notional")
        .join(fx, on=["trade_date", "ccy"], how="left")
        .with_columns((pl.col("notional") * pl.col("rate")).alias("exposure_chf"))
        .group_by("trade_date", "book")
        .agg(pl.col("exposure_chf").sum())
    )

The CPU side was more ordinary: a couple of the scoring routines recomputed the same lookup inside a loop, and one comparison was quadratic in a set that had grown by an order of magnitude since it was written. Neither was clever to fix. They were just hard to see without a profile.

• Measure before changing anything, and keep the profiles alongside the code
• Narrow dtypes on load; most identifier columns did not need 64 bits
• Vectorised expressions over apply, and lazy frames where the pipeline was long enough to benefit
• Regression tests pinned to a stored set of outputs, so an optimisation that changed a result failed loudly

I also looked after the codebase's deployment on Azure, which mostly meant keeping the container image, the pinned dependency set and the scheduled run honest. Separately, I helped a cross-functional team pull measurement data out of custom TIFF files whose tags did not follow the baseline specification; that one ended up as a near-rewrite of the reader, because the original had assumed a single vendor's layout throughout.