Run Batch Embedding Generation at Production Scale

Tier: Core · Workflow area: RAG data pipelines

This playbook builds the batch embedding stage like the data engineering job it is, with the same patterns you already use for any large backfill.

When to reach for this

• •
  Initial index build over a corpus large enough that a naive loop takes hours or breaks API limits.
• •
  Re-embedding the corpus after an embedding-model change (the tax identified in pb-29).
• •
  Recurring incremental embedding as part of CDC-driven index sync (pb-36) that must be safe to retry.
• •
  Choosing between external embedding APIs, Spark UDFs, and warehouse-native functions for a given corpus size.

The workflow

1. 1.

   Make chunks addressable and idempotent first. Key every chunk by a deterministic ID: sha256(doc_id + chunk_index + content_hash). Store chunks in a table (DuckDB locally; your warehouse in production) with an embedded_at and model_version column. Idempotency is a data-model property, not a retry flag.

   Copy

   -- DuckDB / warehouse table driving the pipeline
   CREATE TABLE IF NOT EXISTS chunks (
     chunk_id TEXT PRIMARY KEY, doc_id TEXT, content TEXT,
     content_hash TEXT, model_version TEXT, embedded_at TIMESTAMP
   );
   

2. 2.

   Select work with a query, not a file list. The work queue is simply WHERE embedded_at IS NULL OR model_version <> 'bge-m3-v1'. A crashed run resumes by re-running the same query — nothing to remember, nothing to double-process.

3. 3.

   Batch for the backend you have. Local/self-host path (sentence-transformers on GPU or CPU) is throughput-bound; hosted APIs are rate-limit-bound; warehouse-native (BigQuery AI.GENERATE_EMBEDDING, Snowflake Cortex EMBED_TEXT_1024) moves compute to the data.

   Copy

   import duckdb
   from sentence_transformers import SentenceTransformer
   
   con = duckdb.connect("pipeline.db")
   model = SentenceTransformer("BAAI/bge-m3")
   MODEL_V = "bge-m3-v1"
   
   while True:
       batch = con.execute(
           "SELECT chunk_id, content FROM chunks WHERE embedded_at IS NULL LIMIT 256").fetchall()
       if not batch:
           break
       vecs = model.encode([c for _, c in batch], batch_size=64, normalize_embeddings=True)
       upsert_vectors([b[0] for b in batch], vecs)          # vector store upsert, keyed by chunk_id
       con.executemany(
           "UPDATE chunks SET embedded_at = now(), model_version = ? WHERE chunk_id = ?",
           [(MODEL_V, b[0]) for b in batch])
   

   Hosted equivalent: swap model.encode for the provider's batch embedding endpoint with exponential backoff on 429s; provider batch APIs typically discount ~50% versus synchronous calls and are the right default for million-chunk backfills without a GPU.

4. 4.

   Add the dead-letter path. Chunks that fail after N retries (malformed text, oversize input) go to a chunks_dlq table with the error, and the run continues. The metric that proves this step worked: failure isolation — one poison chunk must not fail the batch, and DLQ rate should be well under 0.5% of chunks.

5. 5.

   Track cost and throughput as run metrics. Log per run: chunks embedded, tokens (estimate via len(text)//4 locally, exact from API usage fields), wall time, and unit cost. Expected output per run, persisted:

   Copy

   run=2026-07-28T10:00 chunks=1_204_331 tokens=402M wall=3h12m dlq=214 (0.02%) cost=$8.04
   

   Warehouse-native comparison: run 10K chunks through AI.GENERATE_EMBEDDING and through the external path, compare unit cost and end-to-end time — this settles the "in-warehouse versus external API" argument with numbers.

6. 6.

   Verify with the re-run test. Immediately re-execute the pipeline. Success criteria, measured: zero chunks re-embedded (work query returns empty), vector-store count equals SELECT count(*) FROM chunks WHERE embedded_at IS NOT NULL, and a spot check of 100 random chunk IDs finds exactly one vector each. Then kill a run mid-flight and re-run: same three checks must still pass.

Verify before you trust it

• •
  Re-run test and kill-and-resume test pass with zero duplicate work (step 6 metrics).
• •
  Vector count reconciles with the chunk table exactly; a mismatch of even 1 means the upsert and the bookkeeping update are not atomic enough — move the embedded_at update after a confirmed upsert.
• •
  DLQ inspected, not just counted: read 10 dead letters and classify causes; systematic causes (encoding, oversize) get fixed upstream, not retried forever.
• •
  Sample 20 embeddings and verify norm ≈ 1.0 if you normalize — a mixed-normalization index degrades retrieval silently.

Pitfalls

• Marking done before the vector store confirms. Update embedded_at only after the upsert succeeds; the reverse ordering loses vectors on crash and the work query will never find them again.
• Rate-limit handling by sleeping a fixed second. Use exponential backoff with jitter honoring retry-after; fixed sleeps either waste hours or synchronize retries into repeated 429 storms.
• One giant transaction. Committing per-batch (256–1,000 chunks) bounds re-work on failure; a single commit at the end means any crash redoes everything.
• Forgetting model_version in the key. A model upgrade must look like "all chunks pending again". Without the version column you will serve a Frankenstein index of mixed embedding spaces — retrieval degrades subtly, not loudly.
• Estimating cost after the run. Compute projected token count and cost from the chunk table before launching a hosted backfill; a surprise five-figure embedding bill is a career metric, not a run metric.

Tools

• sentence-transformers + DuckDBthe fully local pipeline: DuckDB as work queue and bookkeeping, sentence-transformers as the engine; scales further with Spark UDFs on the same pattern.
• Provider batch embedding APIshosted path with ~50% batch discounts; the same work-queue pattern drives request-file creation and result reconciliation.
• BigQuery AI.GENERATE_EMBEDDING / Snowflake Cortex EMBEDwarehouse-native option when the corpus already lives in the warehouse and governance favors keeping it there.

Client-safe variant

For client corpora, the local path (sentence-transformers inside the client VPC or on client-approved compute) avoids any embedding-API data transfer, and the DuckDB-or-warehouse work queue keeps all bookkeeping in client systems. If hosted APIs are approved, still keep the work queue and DLQ client-side, send only chunk text (never document paths or metadata that reveal client identity in a shared workspace), and pin the cost ceiling in writing — batch jobs at client scale are exactly where an unattended pipeline can spend real money. Warehouse-native functions are often the easiest approval path since data never leaves the client's existing cloud contract.