Processing large tables in batches with partitions

Use result.partitions(n) when the work is read-only and can finish in one transaction, and a restartable keyset loop — WHERE id > :last ORDER BY id LIMIT n, committing each batch — whenever the job writes, because committing part-way through a server-side cursor invalidates it. This guide belongs to streaming large result sets with yield_per.

Quick Answer

Streaming and committing do not combine: the cursor belongs to the transaction the commit ends.

Before — committing inside a streamed iteration:

from sqlalchemy import select

from shop.models import Order


async def normalise_orders(session) -> None:
    result = await session.stream_scalars(
        select(Order).order_by(Order.id).execution_options(yield_per=1_000)
    )
    async for order in result:
        order.reference = order.reference.strip().upper()
        await session.commit()        # ends the transaction the cursor lives in
# sqlalchemy.exc.InvalidRequestError / the cursor is gone on the next batch,
# and without the commit this becomes one transaction over the whole table.

After — a restartable keyset loop, one transaction per batch:

import logging

from sqlalchemy import func, select, update
from sqlalchemy.ext.asyncio import async_sessionmaker

from shop.models import Order

log = logging.getLogger("jobs.normalise")


async def normalise_orders(Session: async_sessionmaker, batch_size: int = 1_000) -> int:
    last_id, processed = 0, 0
    while True:
        async with Session.begin() as session:
            ids = list((await session.scalars(
                select(Order.id)
                .where(Order.id > last_id, Order.reference.is_not(None))
                .order_by(Order.id)
                .limit(batch_size)
            )))
            if not ids:
                return processed
            await session.execute(
                update(Order)
                .where(Order.id.in_(ids))
                .values(reference=func.upper(func.trim(Order.reference)))
            )
            last_id = ids[-1]
            processed += len(ids)
        log.info("normalised %d orders (through id %d)", processed, last_id)

Each iteration is its own transaction: it commits, releases its snapshot, and records a position the next iteration continues from. A failure loses at most one batch, and a restart with the last logged id resumes rather than beginning again.

Execution Context & Async Workflow Integration

yield_per puts the result into streaming mode — a server-side cursor on PostgreSQL — and result.partitions(n) yields lists of rows from it. Both belong to the transaction that opened the cursor. commit() or rollback() ends that transaction, and the cursor with it, which is why a batch loop that commits cannot be built on top of streaming.

That leaves two distinct shapes, and choosing between them is the main decision in any large-table job.

Streaming, one transaction. Right when the work is read-only and can complete in one pass: an export, a checksum, feeding rows to an external API. partitions() is the convenient form because most such work batches naturally:

from sqlalchemy import select

from shop.models import Order


async def export_batches(session) -> None:
    result = await session.stream(
        select(Order.id, Order.total_cents)
        .order_by(Order.id)
        .execution_options(yield_per=5_000)
    )
    async for batch in result.partitions(5_000):
        await ship_to_warehouse(batch)      # one call per 5,000 rows

Keyset batches, a transaction each. Required whenever the job writes, and preferable for anything long-running even when it does not, because it is restartable and does not hold a snapshot.

The costs of the single long transaction are worth being concrete about. It holds a snapshot, so VACUUM cannot reclaim rows deleted by other transactions while it runs — a job over a large table for an hour can leave a table measurably bloated. It occupies a pooled connection throughout. It cannot be resumed. And on a hot standby, a long query can be cancelled outright when it conflicts with replay, which is a failure mode that only appears in production.

Two details make keyset batching correct rather than approximately correct. Order by a unique key — the primary key — so id > :last is a total ordering with no ties to skip or repeat; ordering by a timestamp with duplicates loses rows. And re-check the work predicate in each batch (reference IS NOT NULL above), because rows change while the job runs, and a batch selected on stale criteria can act on rows that no longer qualify.

For the write itself, prefer a set-based statement over loading objects. update(Order).where(Order.id.in_(ids)) is one statement per batch; loading a thousand objects and flushing is a thousand UPDATEs unless the ORM can batch them, plus the memory and the identity map. The exception is work that must run per-object Python — validators, events, version counters — which is the trade-off described in using ORM-enabled UPDATE and DELETE statements.

Resolving Warnings, Errors & Common Mistakes

The identity-map growth is worth expanding, because it surprises people who did switch to batches. A session keeps every object it loaded, so a loop that reuses one session for a thousand batches holds a million objects by the end. Opening a session per batch — async with Session.begin() inside the loop, as in the quick answer — solves it completely, and costs one checkout per batch rather than one per row.

Progress reporting deserves to be durable rather than logged. A job that records its position in a table can resume exactly, and its progress is visible to an operator:

from sqlalchemy import select, update

from shop.models import JobProgress


async def resume_point(session, job: str) -> int:
    return await session.scalar(
        select(JobProgress.last_id).where(JobProgress.job == job)
    ) or 0


async def record_progress(session, job: str, last_id: int) -> None:
    await session.execute(
        update(JobProgress).where(JobProgress.job == job).values(last_id=last_id)
    )

Updating that row inside the batch transaction is what makes the resume exact: the work and the position commit together, so the recorded position can never be ahead of the work. Logging it separately, after the commit, leaves a window where a crash re-processes one batch — acceptable when the work is idempotent, and not when it is not.

Throttling is the last practical addition. A batch loop at full speed can saturate the database's write capacity and starve request traffic. A short sleep between batches, or a check of replication lag before continuing, turns the job into a background process rather than an incident:

import asyncio

from sqlalchemy import text

MAX_LAG_SECONDS = 5.0

lag = await session.scalar(text(
    "SELECT COALESCE(EXTRACT(EPOCH FROM (now() - pg_last_xact_replay_timestamp())), 0)"
))
if lag and lag > MAX_LAG_SECONDS:
    await asyncio.sleep(lag)

Advanced: Parallel Batches and Work Claiming

A sequential batch loop is bounded by one connection's throughput. When a job has to finish inside a window, the batches can run concurrently — with two constraints: the concurrency must be bounded by the pool, and no two workers may process the same rows.

For a job that can partition its work statically, the simplest safe split is by key range:

import asyncio

from sqlalchemy import func, select
from sqlalchemy.ext.asyncio import async_sessionmaker

from shop.models import Order


async def parallel_normalise(Session: async_sessionmaker, workers: int = 4) -> int:
    async with Session() as session:
        lo, hi = (await session.execute(
            select(func.min(Order.id), func.max(Order.id))
        )).one()
    if lo is None:
        return 0

    step = (hi - lo) // workers + 1
    ranges = [(lo + i * step, min(lo + (i + 1) * step, hi + 1)) for i in range(workers)]

    async def run(start: int, end: int) -> int:
        return await normalise_range(Session, start, end)

    async with asyncio.TaskGroup() as tg:
        tasks = [tg.create_task(run(start, end)) for start, end in ranges]
    return sum(task.result() for task in tasks)

Each worker owns a disjoint key range, so no coordination is needed, and each runs its own keyset loop inside that range with its own sessions. Keep workers at or below pool_size, for the reasons in running concurrent queries with AsyncSession; key ranges are also uneven when ids are sparse, so expect one worker to finish early.

When the work cannot be partitioned in advance — a queue of rows arriving continuously — claiming is the alternative. SELECT ... FOR UPDATE SKIP LOCKED hands each worker a batch nobody else will touch:

from sqlalchemy import select, update

from shop.models import Job


async def claim_batch(session, size: int = 100) -> list[int]:
    ids = list((await session.scalars(
        select(Job.id)
        .where(Job.status == "pending")
        .order_by(Job.id)
        .limit(size)
        .with_for_update(skip_locked=True)
    )))
    if ids:
        await session.execute(
            update(Job).where(Job.id.in_(ids)).values(status="running")
        )
    return ids

That is the job-queue pattern in full, with its own guide in building a job queue with SELECT FOR UPDATE SKIP LOCKED. For a one-off backfill it is more machinery than key ranges need; for continuous work it is the right shape.

One measurement decides whether any of this is worth it: whether the job is bound by the database or by the application. If each batch's UPDATE is the slow part, more workers make it slower by competing for the same write capacity. If the slow part is per-row Python or an external API call, parallel batches help a great deal.

Choosing a Batch Size

Batch size is the one parameter these jobs always need tuned, and it is tuned against three limits rather than one.

Lock duration. Every batch holds row locks until it commits, and a batch that takes two seconds blocks concurrent writers to those rows for two seconds. For a table serving live traffic, that is the binding constraint: keep batches short enough that a user never waits noticeably. A few hundred to a few thousand rows is the usual range.

Statement efficiency. Very small batches waste round trips: a batch of ten rows spends most of its time on the two statements rather than the work. There is a clear knee in the curve, typically somewhere between a hundred and a few thousand rows, past which larger batches stop helping.

WAL and replication. Each commit produces WAL that replicas must apply. Large batches produce it in bursts, and a job that outruns replication lag either delays replicas or — if a standby is serving reads — starts returning stale data. Watching pg_last_xact_replay_timestamp() between batches, as above, turns that from a surprise into a throttle.

A reasonable procedure is to start at 1,000, measure three things per batch — duration, rows affected, and replication lag after the commit — and adjust:

import time


async def timed_batch(Session, last_id: int, size: int) -> tuple[int, float]:
    started = time.perf_counter()
    async with Session.begin() as session:
        count = await process_one_batch(session, last_id, size)
    return count, time.perf_counter() - started

If batches take under about fifty milliseconds, raise the size; if they exceed a second on a table with live traffic, lower it. Log the numbers rather than inferring them — a job whose batch duration climbs as it runs is usually hitting a missing index on the predicate, not a batch-size problem.

Finally, decide up front what the job does when it fails. A restartable loop with a durable position needs nothing but a retry. A loop that logs its position needs an operator to pass it back. A loop with neither has to start over, which for a table of tens of millions of rows may mean it never completes at all. Writing the position into a table, in the same transaction as the work, is a few lines that turn a fragile script into a job that can be run again without thought — the same reasoning behind resumable migration loops in running Alembic across multiple databases and tenant schemas.

Frequently Asked Questions

Related

• Streaming Large Result Sets with yield_per — The parent guide: server-side cursors and memory.
• Paginating large result sets with keyset pagination — The keyset technique the batch loop is built on.
• Streaming query results as CSV from FastAPI — The read-only streaming case, end to end.
• Using ORM-enabled UPDATE and DELETE statements — Set-based writes for each batch.