Don’t call `.all()` on a training export

A recommender is only as fresh as the last time you could afford to retrain it. When the export step that feeds the trainer starts blowing the memory ceiling on your biggest tenants, “afford” stops being about GPU hours and starts being about whether the job survives at all.

That was the shape of the problem in our training pipeline. Every retrain begins the same way: pull a tenant’s users, items, and clickstream out of Postgres, reshape them into the three files the GeneRec trainer eats — user_metadata.txt, item_catalogue.txt, clickstream.csv — and hand them off. The code that does the reshaping is one class, DataSave. It worked fine for years. Then we onboarded tenants with millions of users, and it stopped working fine.

This is the story of what broke, why the obvious fix wasn’t the fix, and the three changes that actually moved the needle.

The original code fit in your head, which was the problem

Here is the export that shipped for years, lightly trimmed:

def __save_csv(map_query, positive_feedback_query):
    records = map_query.all()
    yield 'userId,itemId,timestamp\n'.encode()
    for row in records:
        output = io.StringIO()
        writer = csv.writer(output)
        writer.writerow([row.user_id, row.job_code, row.start_date or 0])
        yield output.getvalue().encode()

Read it top to bottom and nothing looks wrong. It yields bytes, so it feels like a stream. It even has the word yield in it.

But look at the first line. map_query.all() pulls every row of the clickstream into a Python list before the loop starts. On a tenant with a few thousand users that list is invisible. On a tenant with millions, that single line is the whole ballgame — you’ve materialized the entire result set as SQLAlchemy row objects in RAM, and the container gets OOM-killed before the trainer sees a single byte. The yield downstream is theatre: the expensive thing already happened.

The user-metadata builder had the same tell — records = query.all() — and so did the item builder. Three copies of the same trap.

There’s a second, quieter cost. Every row builds a fresh io.StringIO() and a fresh csv.writer(), formats one line, and throws both away. Ten million rows, ten million throwaway writers. Individually cheap; in aggregate, a tax you pay on the critical path of every retrain.

So there were two problems wearing one coat: unbounded memory (.all()) and per-row object churn (a writer per line). They needed different fixes.

Fix 1: make the stream a real stream

.all() had to go. SQLAlchemy already gives you the tool — yield_per — which fetches rows from the server in bounded batches instead of buffering the whole result. The catch was that the queries reaching DataSave are detached: the session that built them is already closed by the time we want to iterate. So I rebind each query to a fresh, live session that stays open for the whole walk:

def _stream(self, query, batch_size=10_000):
    with acquire_tenant_db_session(self.tenant_id, use_tenant_name=True) as session:
        yield from query.with_session(session).yield_per(batch_size)

That one method is the backbone. Peak memory is now one batch — ten thousand rows — no matter whether the tenant has ten thousand users or ten million. I tried bumping the batch to 50,000 for fewer round-trips, measured it, and reverted: the memory footprint wasn’t worth the marginal throughput. 10,000 was the sweet spot, and I only know that because I looked instead of guessing.

That fixes memory for the metadata files. The clickstream needed something more.

Fix 2: let Postgres format the CSV

The clickstream is the biggest file by far, and it’s the one where per-row Python object construction hurt most. The insight: the database can format CSV faster than Python can, because it does it in C and never builds a Python object at all.

Postgres has COPY ... TO STDOUT WITH (FORMAT csv) for exactly this. It takes a query and streams the rows out already formatted as CSV bytes. No row objects, no csv.writer, no StringIO — the encoding happens server-side and I stream the raw bytes straight through.

The awkward part is the impedance mismatch. psycopg2’s copy_expert uses a push model — it repeatedly calls .write() on whatever sink you hand it — but my caller wants to pull bytes out of a generator. Those two don’t compose directly. A SpooledTemporaryFile is the bridge: I let COPY push the whole result into the spool, then pull it back out in 1 MB chunks and yield those.

with SpooledTemporaryFile(max_size=64 * 1024 * 1024, mode="w+b") as sink:
    cursor.copy_expert(copy_sql, sink)   # push: Postgres writes into the spool
    sink.seek(0)
    while chunk := sink.read(1024 * 1024):
        yield chunk                       # pull: we stream it back out

The max_size is the trick that keeps this honest. A SpooledTemporaryFile lives in memory up to the threshold and silently spills to disk beyond it. So the common case (a few MB) never touches the filesystem, and the pathological case (a huge tenant) degrades to a temp file on disk instead of an OOM. Bounded memory, no branching in my code — the standard library does the deciding.

Two details in the COPY path cost me real time. Both come from the same root: COPY takes a SQL string, but I’m handed a SQLAlchemy query object. Bridging those two is where the sharp edges live.

Why I can’t just stringify the query

The obvious move is to render the query to a literal SQL string and paste it inside COPY (...). SQLAlchemy even offers this — compile(compile_kwargs={"literal_binds": True}) inlines the parameter values directly into the SQL text.

It works right up until a datetime shows up in a WHERE clause. literal_binds only knows how to inline the types it has a literal renderer for — numbers, strings, booleans. A Python datetime has no literal SQL form that SQLAlchemy will emit, so it raises rather than guess. And my query filters on timestamps, so literal_binds was a dead end from the start.

So I split the job in two. First, compile the query to a parameterised statement — SQL text with placeholders, plus a separate bag of parameter values — instead of trying to inline anything:

compiled = map_query.with_session(session).statement.compile(
    dialect=postgresql.dialect(),
    compile_kwargs={"render_postcompile": True},
)

render_postcompile matters here: things like IN (...) clauses expand to the right number of placeholders at this step, so the parameter bag and the placeholders line up exactly.

Then I hand the text and the values to psycopg2 and let it do the inlining, via mogrify:

inner_sql = cursor.mogrify(str(compiled), compiled.params)
if isinstance(inner_sql, bytes):
    inner_sql = inner_sql.decode()

mogrify is the driver’s own “produce the exact SQL I would send” call. It’s the layer that actually talks to Postgres, so it knows how to render every type Postgres accepts — including a datetime — and it does the quoting and escaping correctly, which means no SQL-injection hole even though I’m building a string. (It hands back bytes, hence the decode.) The division of labour is the whole trick: SQLAlchemy builds the query shape; psycopg2 binds the values. Neither one can do the other’s half.

Now inner_sql is a complete, self-contained SELECT with every value already baked in — safe to drop inside a COPY (...) subquery, because COPY can’t take bound parameters of its own.

Why the COPY query says `COALESCE(start_date, 0)`

The second detail is smaller and easier to get wrong precisely because it’s small.

The old Python path formatted each timestamp like this:

value = getattr(row, 'start_date', "") or 0

Read what that actually does with a NULL. Postgres hands back None; getattr returns it (the column exists, so the "" default never fires); then None or 0 evaluates to 0. So a missing timestamp landed in the CSV as the integer 0. Not empty string, not NULL, not None — 0. Downstream consumers of that CSV have been reading 0 for missing timestamps for as long as this code has run.

When the formatting moved from Python into the database, that coercion had to move with it. Postgres won’t apply Python’s truthiness rules for me — a NULL column comes out of COPY as an empty field, which is a different value than 0. So I made the SQL do explicitly what the Python did implicitly:

COPY (SELECT user_id, job_code, COALESCE(start_date, 0) AS start_date
      FROM (<inner_sql>) AS _mq) TO STDOUT WITH (FORMAT csv)

COALESCE(start_date, 0) is the SQL spelling of start_date or 0: hand back start_date unless it’s NULL, in which case hand back 0. Same output byte-for-byte.

This is the part of a “performance” change that has nothing to do with performance and everything to do with trust. A rewrite that’s 10× faster but silently turns 0 into an empty field isn’t an optimization — it’s a data bug with good benchmarks. The edge case that looks too trivial to mention is exactly the one that slips through, because nobody writes a test for the timestamp that was never there. The COALESCE is one function call; noticing it was needed was the actual work.

Fix 3: batch the rows Postgres can’t format for me

Not every clickstream row can go through COPY. The positive-feedback rows do a jobs.split('/') — one database row fans out into several CSV rows — so there’s no single SELECT that produces them. Those still have to be formatted in Python.

But “in Python” doesn’t have to mean “a fresh writer per row.” I build one StringIO and one csv.writer, write into the shared buffer, and only flush — encode and yield — every 10,000 rows:

buffer = io.StringIO()
writer = csv.writer(buffer)
n = 0
for row in self._stream(positive_feedback_query):
    for idx, job_code in enumerate(row.jobs.split('/')):
        writer.writerow([f"{row.user_id}-posFeedback", job_code, idx])
        n += 1
        if n >= _CSV_YIELD_BATCH:
            yield buffer.getvalue().encode()
            buffer.seek(0); buffer.truncate(0); n = 0
if n:
    yield buffer.getvalue().encode()

One encode() per batch instead of per row, one writer for the whole export, memory bounded to roughly one batch of text. Same CSV output, a fraction of the object churn.

The trap in the batching, and why the metadata files don’t get it

Here’s where the two files diverge, and it’s the part I’d most want a future maintainer to internalize.

The clickstream can batch as many rows per yielded chunk as it likes, because the trainer writes those bytes to a file raw — it doesn’t care where the chunk boundaries fall.

The metadata files cannot. Each yielded chunk of user_metadata.txt is one str(dict) record, and downstream the uploader pickles each yielded item individually, then the trainer reads it back and runs ast.literal_eval() on it. And ast.literal_eval parses exactly one Python expression. The moment you “optimize” the metadata builder by joining several records into one yielded chunk, the trainer literal-evals a two-record blob and dies with a SyntaxError on the second one.

So the contract is: clickstream batches, metadata yields one record per chunk. Same file, opposite rule, and the reason is three services away in code you’re not looking at. The kind of invariant that’s invisible until you break it — so it lives in a comment now, right above the loop.

A one-character bug worth its own paragraph

While I was in there, I hoisted a loop-invariant out of the per-user loop. The metadata builder branches on which columns are categorical vs. text, and one branch depends on whether a guid column exists. The old code recomputed the subquery’s column collection on every single row. I lifted it out once:

subquery_cols = list(query.subquery().c)
has_guid = any(c.name == 'guid' for c in subquery_cols)

The has_guid line hides a real bug I nearly reintroduced. subquery_cols is a plain Python list of Column objects. Write the natural-looking 'guid' in subquery_cols and it compares the string 'guid' against Column objects — always False, forever. SQLAlchemy’s ColumnCollection overloads __contains__ to match by name, so the old ... in query.subquery().c worked; a plain list does not. You have to compare c.name. It’s one word of difference and it silently drops the entire competency feature from the export. Tests caught it. Comments now guard it.

What actually moved the needle

Three changes, in order of impact:

.all() → yield_per. This is the one that turns “crashes on big tenants” into “runs on big tenants.” Everything else is speed; this is survival.
Server-side COPY for the clickstream. Postgres formats CSV in C and hands you bytes. Let it.
Batch the Python rows you can’t hand to Postgres. One writer, flush on an interval.

The lesson underneath all three: yield in a function does not make it lazy. A generator that opens with .all() is an eager loader wearing a lazy costume, and the only way to know the difference is to ask where the memory peaks — not where the yield keyword is.

And one meta-lesson I keep relearning: I didn’t touch the ~200s of setup cost that runs before any of this streams, even though it’s the single biggest number in the export. I measured it, found it was a one-time-per-export cost buying a correctness-critical decision about which columns ship, and left it alone. The streaming rewrite captured the value that was actually free to capture. Knowing which slow thing to leave slow is half the job.

Don’t call .all() on a training export