Where the memory went in a Pandas pipeline

A batch job I looked after at a bank read a day of position data, joined it against a slower-moving reference set, and wrote the aggregates into Postgres for a risk model to consume the next morning. It had been comfortable for years. Then the input roughly doubled and the job started being killed for exceeding its memory limit, always at the same stage, always after most of an hour of work.

The reflex is to ask for a larger machine, and we did get one, and it bought about a month. What actually fixed it was an afternoon spent finding out where the memory was going, which was not where any of us would have guessed.

Measuring before changing

The cheapest useful measurement is df.info(memory_usage="deep") on each frame as it is created. Deep is the important part: without it, an object column reports the size of its pointers rather than the size of the strings they point at, and that is precisely the case where the number matters.

Two things fell out of that. Most of the resident memory was held by four string columns with a few dozen distinct values between them, all stored as object. And the peak — which is the number the scheduler cares about — was not the size of any single frame but the moment during a merge when both inputs and the result existed at once.

• Peak memory is what gets you killed; steady-state size is only a proxy for it.
• Low-cardinality strings as category were the largest single win, and cost one line.
• Identifier columns were int64 by default when int32 was plainly enough.
• Every intermediate frame still bound to a name in an outer scope is memory you have chosen to hold.

Aggregating per chunk

The structural change was to stop building one large frame and then reducing it. Reading in chunks and reducing each chunk immediately means the raw rows can be freed as soon as they have been counted, and the thing that accumulates is the aggregate, which is small.

import pandas as pd

DTYPES = {
    "account_id": "int32",
    "product": "category",
    "currency": "category",
    "exposure": "float32",
}

def load_exposure(path: str) -> pd.DataFrame:
    partials = []
    reader = pd.read_csv(path, usecols=list(DTYPES), dtype=DTYPES, chunksize=500_000)
    for chunk in reader:
        # Reduce inside the loop so the raw chunk can be collected.
        partials.append(
            chunk.groupby(["product", "currency"], observed=True)["exposure"].sum()
        )
    return pd.concat(partials).groupby(level=[0, 1]).sum().reset_index()

The float32 deserves a caveat. Halving the width of a float column is only defensible if you know the precision the output needs, and for a figure reported to the nearest thousand it plainly was. On a column feeding a discounting calculation I would have left it as it was and found the memory elsewhere.

A frame that fits in memory is not the goal. A pipeline whose peak does not grow with the input is.

The obvious next step is Polars, and for a rewrite I would start there, because the lazy API does the projection and predicate pushdown that I ended up doing by hand. But these changes took an afternoon and a rewrite would have taken a fortnight, and the job now runs on a smaller instance than it did before the input doubled. That is the ordering I would keep: measure, make the cheap structural change, then consider changing tools.