fix(knowledge): read the projection-filled probe off the partial index instead of a sequential scan (#8139)

* fix(knowledge): read the projection-filled probe off the partial index instead of a sequential scan

* test(knowledge): pin the backfill probe's ordered, capped shape

Author: waleedlatif1

apps/sim/lib/knowledge/search/projection-source-acl-backfill.test.ts

@@ -78,12 +78,21 @@ describe('runProjectionSourceAclBackfill', () => {
 
   it('analyzes and warms the projections on the same connection once both are filled, before closing it', async () => {
     await runProjectionSourceAclBackfill({})
-    /** A row whose document is gone is not the fill's to finish; the probe joins the document. */
-    expect(
-      mockUnsafe.mock.calls.some(([query]) =>
-        String(query).includes('JOIN document d ON d.id = s.document_id WHERE s.acl IS NULL')
+    /**
+     * A row whose document is gone is not the fill's to finish; the probe joins the document. It is
+     * ordered and capped so only the unfilled-rows index can serve it: an `EXISTS` drops both and
+     * leaves the planner a sequential scan of the projection.
+     */
+    const probes = mockUnsafe.mock.calls
+      .map(([query]) => String(query).replace(/\s+/g, ' '))
+      .filter((query) => query.includes('AS unfilled'))
+    expect(probes).toHaveLength(2)
+    for (const probe of probes) {
+      expect(probe).not.toContain('EXISTS')
+      expect(probe).toContain(
+        'JOIN document d ON d.id = s.document_id WHERE s.acl IS NULL ORDER BY s.id DESC LIMIT 1 ) IS NOT NULL AS unfilled'
       )
-    ).toBe(true)
+    }
     expect(mockUnsafe.mock.calls.map(([query]) => query)).toEqual(
       expect.arrayContaining(['ANALYZE embedding_search', 'ANALYZE embedding_keyword_tin'])
     )

apps/sim/lib/knowledge/search/projection-source-acl-backfill.ts

@@ -148,14 +148,17 @@ export async function runProjectionSourceAclBackfill(
 /**
  * Whether no projection still holds a row the fill could give its source and ACL: a row without
  * them whose document exists. A row whose document is gone is not the fill's to finish and never
- * counts as left. Each read is one index probe while any such row remains.
+ * counts as left. Each read is one probe of the unfilled-rows index while any such row remains:
+ * ordered by id and capped at one row so the planner cannot take a sequential scan, which an
+ * `EXISTS` would leave open by dropping the order and the limit.
  */
 async function projectionsFilled(sql: postgres.Sql): Promise<boolean> {
   for (const projection of PROJECTION_SOURCE_ACL_TABLES) {
     const [row] = await sql.unsafe<Array<{ unfilled: boolean }>>(
-      `SELECT EXISTS (
-        SELECT 1 FROM ${projection} s JOIN document d ON d.id = s.document_id WHERE s.acl IS NULL
-      ) AS unfilled`
+      `SELECT (
+        SELECT s.id FROM ${projection} s JOIN document d ON d.id = s.document_id WHERE s.acl IS NULL
+        ORDER BY s.id DESC LIMIT 1
+      ) IS NOT NULL AS unfilled`
     )
     if (row?.unfilled) return false
   }

apps/sim/lib/knowledge/search/queries.test.ts

@@ -2822,6 +2822,24 @@ describe('filters on a resolved scope', () => {
     expect(caps.at(-1)).toBe('20000')
   })
 
+  it('looks for an unfilled row through the ordered, capped read the partial index serves', async () => {
+    traversedRows = [{ id: 'a' }]
+    rerankRows = [hit('a', 'src-a')]
+    queueTableRows(schemaMock.embedding, rerankRows)
+    await handleVectorOnlySearch({
+      ...params,
+      permitted: { kind: 'unbounded', broad: true },
+      accessPlan: plan(),
+    })
+    const probes = statements().filter((query) => query.sql.includes('AS unfilled'))
+    expect(probes).toHaveLength(1)
+    /** An `EXISTS` drops its order and limit, and the planner then takes a sequential scan. */
+    expect(probes[0].sql).not.toContain('EXISTS')
+    expect(probes[0].sql.replace(/\s+/g, ' ')).toContain(
+      'SELECT ( SELECT ? FROM ? WHERE ? IS NULL ORDER BY ? DESC LIMIT 1 ) IS NOT NULL AS unfilled'
+    )
+  })
+
   it('tests the date through the document inside an on-row walk when the filtered set is unbounded', async () => {
     traversedRows = [{ id: 'a' }]
     rerankRows = [hit('a', 'src-a')]

apps/sim/lib/knowledge/search/queries.ts

@@ -115,7 +115,13 @@ const PROJECTION_FILLED_TTL_MS = 60_000
 
 /**
  * Whether the ranking projection still holds rows the backfill has not filled. Read off the
- * unfilled-rows index in microseconds and remembered briefly: the answer only ever changes once.
+ * unfilled-rows index in milliseconds and remembered briefly: the answer only ever changes once.
+ *
+ * The read asks for the last unfilled row by id, not whether one exists: an `EXISTS` drops its
+ * order and limit, and while most rows are unfilled the planner expects a sequential scan to
+ * meet one at once, then walks the whole projection when the unfilled rows sit past the filled
+ * ones. Ordered by id and capped at one row, the read can only be the partial index, whose
+ * last entry is the row the fill reaches last.
  */
 const projectionFilled = new LRUCache<
   ProjectionSourceAclTable,
@@ -134,7 +140,10 @@ const projectionFilled = new LRUCache<
     try {
       const [row] = await runSearchQuery(context.budget, context.stage, (executor) =>
         executor.execute<{ unfilled: boolean }>(sql`
-        SELECT EXISTS (SELECT 1 FROM ${table} WHERE ${table.acl} IS NULL) AS unfilled`)
+        SELECT (
+          SELECT ${table.id} FROM ${table} WHERE ${table.acl} IS NULL
+          ORDER BY ${table.id} DESC LIMIT 1
+        ) IS NOT NULL AS unfilled`)
       )
       return !row?.unfilled
     } catch {