fix(firestore-vector-search): restore batched backfill - #3097
Conversation
There was a problem hiding this comment.
Code Review
This pull request restores the batched backfill and update triggers for the Firestore Vector Search kit, processing documents in chunks via Cloud Tasks and tracking progress in a metadata document. The review feedback highlights several critical issues in the backfill implementation, including race conditions when enqueuing the first task and updating progress counters outside of a transaction, a transaction retry side-effect that can duplicate documents in arrays, and a potential parsing bug in task ID generation if the instance ID contains the substring 'task-'.
|
Review of the changes in this PR against the extension. Three problems, all in the path this PR touches. 1. The OpenAI client's batch size is 1, so
2. The enqueue documents are committed after the first task is dispatched, so the first task can race its own successor.
3. The The README enumerates the internal Firestore paths this kit owns, and the restored task thread adds |
|
Second review pass. One problem, introduced by this PR. With
The extension never hit this because it ran the backfill on install and the update on configure, never both at once, so a single metadata document was enough.
Also worth stating plainly in the description, since it is a difference a reviewer comparing against the extension will look for: both passes record the full metadata set. The extension's |
|
Third review pass. Two small problems, both in the surface this PR adds; no behaviour differences left against the extension. 1.
2. "Delete That document is also the parent of the Also folded in: |
d8f14f1 to
aa9954b
Compare
The backfill and update triggers loaded the whole collection with a single `get()` and enqueued one Cloud Task per document, each embedding one document with one API call, and neither pass consulted the index metadata document, so every redeploy re-embedded the entire collection. Restore the extension's task-thread backfill: enumerate the collection with `listDocuments()`, chunk it into 50 document ids per task, run one task at a time with each task enqueueing its successor, and embed each chunk in provider-sized batches with a single API call per batch. Documents are marked `BACKFILLED` or `FAILED_BACKFILL` with a `completeTime` in one batched write, and a failed batch marks its documents failed instead of failing the task. Both passes are gated on `_<instance id>/index`, comparing the embedding provider, the vector dimension and the input/output field names against the last recorded pass. Unlike the extension, the progress counters merge into that document rather than replacing it, so the gate survives its own first pass, and the enqueue-document batch commits its remainder so a thread of more than 50 chunks does not stall on a task that was never recorded.
…ks first Three problems found reviewing the restored backfill. `OpenAiEmbedClient` declared a batch size of 1 where the extension declared 16, so the chunk runner made one embedding request per document for the `openai` provider — the cost this restoration was meant to remove. `enqueueTaskThread` dispatched the first task before committing the enqueue documents, so for any thread of 2 to 50 chunks the first task could finish and look for a successor that had not been written yet. Every chunk is now recorded before anything is dispatched. The `enqueues` subcollection was missing from the README's list of the internal Firestore paths this kit owns.
`handleInit` enqueued both triggers when `DO_BACKFILL` and `UPDATE_ON_CONFIGURE` were both set. Now that both passes share the index metadata document as their task thread and mint the same task ids, two concurrent threads overwrite each other's enqueue documents and interleave their counter increments, so neither reaches its total. Enqueue the update trigger only when the backfill trigger was not enqueued. The backfill pass is a strict superset of the update pass, so nothing is lost, and either setting on its own is unaffected.
`BackfillProcess` is part of the `./lib` surface, and every one of its signatures referenced a `DocumentData` alias that was not exported and shared its name with Firestore's own type, so a consumer implementing the interface could not name its parameters. Renamed to `BackfillDocumentData` and exported. Also add `"PROCESSING"` to `BackfillJobStatus`, the one status the module wrote without it, and make the README's "force a re-embed" advice say that deleting the metadata document leaves its `queries` subcollection alone.
…ansaction
Two findings from the automated review that cost nothing in parity terms.
`getNextTaskId` took the counter with `split("task-")[1]`, which reads the
wrong segment when the instance id itself contains that substring, so the
parse yielded NaN. It now captures the digits from the same anchored pattern.
`getValidDocs` accumulated into arrays declared outside `runTransaction`.
Firestore may run the callback more than once, and a retry would have appended
the same documents again, overstating the chunk's counts. They are now
collected inside the callback and returned from it.
…ass rule in the queue-target test The queue-target test landed on kits after this branch was cut. It stubbed a Firestore whose collection get() returns one document, and expected init to enqueue both trigger queues when DO_BACKFILL and UPDATE_ON_CONFIGURE are set. Both assumptions are now wrong: a trigger pass reads and writes the index metadata document, enumerates the collection with listDocuments(), and records the enqueue documents in a batch before dispatching; and init enqueues only the backfill trigger when both passes are on, because the two passes share one task thread and the backfill covers every document the update pass would. The stub gains doc(), listDocuments() and batch(), the init case asserts the backfill trigger alone, and a second case pins the update trigger on its own.
cb1a886 to
395d746
Compare
|
Rebased onto What changed in the rebase:
Verified on the rebased head: 159 tests pass (113 on |
cabljac
left a comment
There was a problem hiding this comment.
Rebased this onto kits earlier today (see the comment above for what changed). The code side looked good to me and I have no code changes to ask for; the three inline comments are all about the README describing what the code actually does, since the Differences section is where a migrating user will look.
Fair warning: the parity comparison against the extension and the helper package was AI-led, and I have only spot-checked it. If any of the extension behaviour I describe does not match what you saw when you ported it, push back and I'll go and look properly.
| gated on `DO_BACKFILL` instead. The kit reads `UPDATE_ON_CONFIGURE`, so the two | ||
| passes are controlled independently — `DO_BACKFILL` after the first deploy, | ||
| `UPDATE_ON_CONFIGURE` after every redeploy. |
There was a problem hiding this comment.
I think this sentence doesn't match what the code does. initVectorSearch runs on both afterFirstDeploy and afterRedeploy (src/index.ts:118-123), and handleInit returns after enqueuing the backfill whenever doBackfill is on (src/handlers.ts:207-212), so with DO_BACKFILL=true it is the backfill trigger that gets enqueued on every redeploy too, and UPDATE_ON_CONFIGURE never gets a look in. Lines 301-302 already say "after your first deploy and after every redeploy", so the two paragraphs contradict each other.
That also makes the one-pass rule a deviation worth naming. The extension maps install to backfillTrigger and configure to updateTrigger (extension.yaml:400-406), and the update pass only touches documents that already have an embedding, so a reconfigure that changed INPUT_FIELD_NAME left never-embedded documents alone. Here it backfills them as well. Probably the right call given the shared thread, but it changes which documents get written on a redeploy, so it belongs in the Differences list rather than only in the handleInit comment.
One more thing this section should say, because "redeploying without changing anything enqueues nothing" is only true from the second kit deploy. On a migrated instance the first deploy re-embeds the whole collection: the extension's _<instance id>/index document holds only the progress counters after any pass that enqueued tasks (its set() replaced the config fields), so the gate opens, and the pass then treats every extension-written document as unprocessed because it reads status.state (src/backfill.ts:470-472) where the extension wrote status.<instance id>.state. That is a bigger bill than the extension's reconfigure, which skipped documents the write trigger had marked COMPLETED and re-embedded only the ones its backfill had marked BACKFILLED. With DO_BACKFILL=true it lands on the first kit deploy, so people should expect it. Worth checking my reading of the status path there, since that is the #3011 shape and I may be behind on where that landed.
Could you reword this paragraph to what actually happens (both settings act on every deploy, backfill wins when both are on), and add the install/configure mapping and the first-deploy re-embed to the Differences list?
|
|
||
| The extension's gate did not survive its own first pass, because the progress | ||
| counters it wrote to the same document replaced the recorded configuration. The | ||
| kit merges instead, so the comparison fields persist and the gate holds on every | ||
| later deploy. To force a full re-embed without changing any setting, delete the | ||
| `_<instance id>/index` document; its `queries` subcollection is untouched, so |
There was a problem hiding this comment.
This section names two deviations (the merge, and the one-pass rule), but reading backfill.ts against the helper package I count more, and the rule is that every one is documented here so a reader can decide whether it matters to them. The ones I can see:
- Every enqueue document is committed before the first task is dispatched, and the remainder is always committed (
backfill.ts:185-216). The extension dispatched task-1 before its own enqueue document was in the batch (trigger.ts:70-83) and never committed the chunks past the last multiple of 50 (trigger.ts:89), so 51 to 99 chunks lost the tail. - Update-pass failures stay on their own document (
backfill.ts:431-438,392-398). The extension dropped rejected results and shifted the rest, so one document could receive another's embedding. - A batch result shorter than its input marks the trailing documents
FAILED_BACKFILL. The extension marked themBACKFILLEDwith no new vector. Side effect of the fix above, but observable. - A document id with no document (a subcollection parent) is skipped (
backfill.ts:462). The extension'sshouldBackfillthrew aTypeErrorthere and the task retried until the queue gave up. - The update pass records all four comparison fields on the metadata document (
handlers.ts:324-333). The extension recorded onlyembeddingProviderfor that pass, so even without theset()wipe the other three comparisons would have readundefinedand its gate would have opened on every reconfigure. getValidDocscollects inside the transaction callback (backfill.ts:453-484). The extension accumulated outside it, so a transaction retry double-counted.
All of these look justified to me under "matching would reproduce a bug", and the PR body already argues most of them. Could you add each to this list with one line on what the extension did? The handleInit comment and the PR body are not where a user will look.
| if (ctx.config.doBackfill) { | ||
| await enqueueBackfillTrigger(ctx); | ||
| // The two passes share one task thread on the index metadata document, and | ||
| // the backfill pass covers every document the update pass would. Running | ||
| // both at once would have them overwrite each other's progress. | ||
| return; | ||
| } | ||
| if (ctx.config.updateOnConfigure) { | ||
| await enqueueUpdateTrigger(ctx); |
There was a problem hiding this comment.
One deviation here that I think is unclaimed, and it only half holds. The extension's install always ran backfillTrigger, and updateOrCreateMetadataDoc created the document with the full config before checking DO_BACKFILL (metadata_document.ts:64-69, called from firestore_process_backfill_trigger.ts:14-23), so with DO_BACKFILL=false the document still existed. Flipping the setting to true later and reconfiguring then found a matching document and did not backfill. Here nothing is enqueued when both settings are off, so no document is written and flipping DO_BACKFILL on later does run the pass. With UPDATE_ON_CONFIGURE=true though, the update trigger writes the document on the first deploy and the flip is gated out exactly as on the extension. Worth a sentence in the Differences list either way, and the "delete the metadata document to force a pass" advice covers the trap. Worth double-checking my reading of the extension's ordering there.
Closes #3012. Parity ledger: #2974,
firestore-vector-search§9.The kit's backfill and update triggers loaded the whole collection with one
collection().get()and enqueued one Cloud Task per document, each embedding a single document with a single API call. Neither pass consulted the index metadata document, soDO_BACKFILL/UPDATE_ON_CONFIGUREre-embedded the entire collection on every deploy.This restores the extension's task-thread backfill (§9a–9g).
Changes
src/backfill.ts: the index metadata gate, the task-thread enqueue, the dispatched-task runner, the chunk runner and the eligibility rules, ported from the extension's@invertase/firebase-extension-utilitiesdistributed-task helpers.listDocuments()instead ofcollection().get(), so the trigger no longer holds every document in memory (§9a)._<instance id>/index(backfillJobsTotal/Processed/Skipped/Failed,backfillStatus) with one enqueue document per chunk under_<instance id>/index/enqueues(§9b).EmbedClient.batchSize) with onegetEmbeddings()call per batch; the update pass embeds one document per call at a chunk size of 50, matching the extension's two processes (§9c).embeddingProvider,dimension,inputFieldandoutputFieldagainst the last recorded pass, and skipping when none of them changed (§9d).BACKFILLEDis skipped. The update pass additionally requires an existing embedding (§9e).BACKFILLEDorFAILED_BACKFILLwith acompleteTime, rather than a per-documentset(merge)ofCOMPLETED(§9f).README.mdandCHANGELOG.mdupdated: the "backfill is one task per document" and "UPDATE_ON_CONFIGUREre-embeds on every deploy" sections are replaced, the status-field section now recordsBACKFILLED/FAILED_BACKFILL, and the batched behaviour moves into "Unchanged".tests/backfill.test.ts: 38 tests over an in-memory Firestore covering the gate, the chunking and dispatch, every skip rule, the batch and single-document paths, both failure paths, and the trigger gating.Three deliberate deviations from the extension
Each one is a case where porting the extension verbatim would leave the behaviour this issue asks for unreachable.
taskThreadTriggerwrites the progress counters to the metadata document withset()and no merge, which wipes the four comparison fields on the first pass. Every later deploy then reads a document without them, the comparison fails, and the gate never holds — the exact symptom in this issue. The kit merges, so the gate survives its own first pass.counter % batchSize === 0 || chunks.length < batchSize), so a thread of more than 50 chunks (over 2,500 documents) leaves its trailing enqueue documents uncommitted, and the task before them fails permanently withNext task document … does not exist. The kit commits the remainder after the loop.batchFn, soProcess.batchProcessmaps each document through the single-document path and drops rejected results, shifting every later result one slot and writing one document's embedding onto another. The kit keeps the slots aligned and marks only the documents that actually failed.Two smaller notes: a document id returned by
listDocuments()that has no document (a phantom parent of a subcollection) is skipped rather than crashing the task, which is what the extension'sshouldBackfillintended before it dereferencedundefined; and there is nosetProcessingStateequivalent for a kit, so the extension's install-UI progress reporting is logs plus the progress document instead — already a Notes item on #2974.Out of scope and unchanged: the status field's flat shape (#3011 / §6) and the lifecycle wiring that lets
afterRedeployreach both passes (§10). Task ids arekit-<instance id>-task-Nrather than the extension'sext-prefix, matching the kit's queue naming; they are internal to the progress document.Rebased onto
kitsafter finding the branch point was 51 commits stale; the only overlap wastests/embeddings.test.ts(the custom-endpoint suite from #3042), kept alongside the new OpenAI suite. The task functions now carry the provider secrets from #2987, which the backfill task needs.Testing
Automated review raised four points, all against the first push. One was already fixed by a later commit, two are fixed in cb1a886 (task-id counter now captured rather than split on
task-, so an instance id containing that substring parses;getValidDocscollects inside the transaction callback so a retry cannot duplicate entries), and one was declined as a parity reduction: moving the progress counters into a transaction would replace the extension'sFieldValue.incrementwrites with absolute ones, and duplicate Cloud Tasks delivery overshoots upstream too.Unit —
vitest runinkits/firestore-vector-search, 101 tests, all passing (was 49).tests/backfill.test.tsadds 45 over an in-memory Firestore that records every write and refusescollection().get(), plus 4 OpenAI client tests intests/embeddings.test.ts:embeddingProvider/dimension/inputField/outputFieldchanged; and — the regression this issue is about — still gates after the progress counters have been written to the same document.getEmbeddingsonce per provider-sized batch (a 3-id chunk at batch size 2 becomes[["one","two"],["three"]]), single-eligible-document chunks take thegetSingleEmbeddingpath, and both markBACKFILLEDin one committed batch.COMPLETEDorERRORare all skipped;BACKFILLEDis re-embedded; the update pass additionally skips documents with no existing embedding.FAILED_BACKFILLand the task still resolves; a rejected single document does the same; on the update pass only the document that actually failed is marked, and the one that succeeded still gets its embedding.DONEinstead of dispatching once every job is accounted for, a progress document without counters throwsInvalid task document, and a missing successor throws.listDocuments(), gate on unchanged configuration, enqueue nothing for an empty collection, and swallow an enqueue failure.handleInitenqueues exactly one trigger when both settings are on.Deployed — repacked the kit into a throwaway codebase in the test project (
fvs-pr3097: its own source, config, collectionfxkits_pr3097and secrets,EMBEDDING_PROVIDER=vertex), seeded 60 documents before any function existed, then deployed. Everything below is from the live run; that codebase, its functions, secrets and data have since been deleted.Seed: 55 plain documents, one already
BACKFILLED, oneCOMPLETED, oneERROR, one with a non-string input, one with an empty-string input.afterFirstDeployraninitVectorSearch, which created the vector index and enqueued the backfill trigger. The trigger loggedNo existing metadata doc found->Creating a new metadata doc->Found 60 documents in the collection fxkits_pr3097->Enqueuing backfill tasks->Committing the batch...->Enqueuing the first task kit-fvs-pr3097-task-1->2 tasks enqueued successfully. The commit precedes the dispatch, which is the ordering fix.task-1carrying a 50-id chunk andtask-2carrying 10, both endingDONE.task-1loggedSuccessfully enqueued task kit-fvs-pr3097-task-2, so the thread chained instead of running both at once.BACKFILLED, every one with acompleteTimeand a nativeVectorValueof 768 dimensions. The already-BACKFILLEDdocument was re-embedded, as the extension did.skip-completedandskip-error:is not in the correct state to be backfilled;skip-nonstringandskip-empty:is not valid for fvs-pr3097 process. All four were left untouched, with no embedding and their original status intact.{backfillJobsTotal: 60, backfillJobsProcessed: 56, backfillJobsSkipped: 4, backfillJobsFailed: 0, backfillStatus: "DONE"}, matchingCurrent state: 56 processed, 4 skipped, 0 failed out of 60 total tasksin the logs.embeddingProvider,dimension,inputField,outputFieldandcollectionNamealongside those counters. With the extension's replacingset()the counters would have wiped all five, which is what made its gate a no-op.Embedding configuration is unchanged for fxkits_pr3097, no pass required.Nothing enumerated, nothing enqueued, no embedding calls. This is the case fix(firestore-vector-search): restore batched backfill #3012 reports, and it now costs nothing.DO_BACKFILL=false/UPDATE_ON_CONFIGURE=true, cleared the metadata document and redeployed. It ran throughupdateTaskover 61 documents in chunks of 50 and 11: 56 processed, 5 skipped, 0 failed,DONE. Its extra rule shows up as a different log reason for the same documents:skip-completedandskip-errorare nowis not valid for ... process(no embedding to update), while theCOMPLETEDdocument that does have one isis not in the correct state.embedOnWriteon a fresh document still produced an embedding andstatus.state: "COMPLETED".queryOnWritenot exercised — the vector index on the brand-new collection was still building (FAILED_PRECONDITION: That index is currently building). Unrelated to this change; the query paths are untouched here and covered by the existing unit suites.The same gate behaviour was also confirmed on the pre-existing
firestore-vector-search-a91f6c2einstance: metadata document created on first deploy, gate held across two unchanged redeploys, and gate reopened (Updating existing metadata doc) whenINPUT_FIELD_NAMEwas changed.