v0.8.55: projection index for sim search

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 {

packages/db/script-migrations/0021_embedding_search_connector.test.ts

@@ -1,14 +1,141 @@
 /**
  * @vitest-environment node
  */
-import { backfillProjectionSourceAcl } from '@sim/db/script-migrations/0021_embedding_search_connector'
+import {
+  backfillProjectionSourceAcl,
+  PROJECTION_SOURCE_ACL_PAGE_RETRIES,
+} from '@sim/db/script-migrations/0021_embedding_search_connector'
 import type { Sql } from 'postgres'
-import { describe, expect, it, vi } from 'vitest'
+import { afterEach, describe, expect, it, vi } from 'vitest'
 
 /** A session that must never be reached: every case below is refused before the first page. */
 const untouched = { begin: vi.fn() } as unknown as Sql
 
+type PageRow = { scanned: number; filled: number; last_id: string | null }
+
+/** The message the database pairs with each cancellation SQLSTATE. */
+const CANCELLATION_MESSAGES: Record<string, string> = {
+  '55P03': 'canceling statement due to lock timeout',
+  '57014': 'canceling statement due to statement timeout',
+}
+
+/** A driver error carrying a SQLSTATE, the shape `postgres` throws. */
+function postgresError(code: string, message = CANCELLATION_MESSAGES[code] ?? 'failed'): Error {
+  return Object.assign(new Error(`${message} (SQLSTATE ${code})`), { code })
+}
+
+/**
+ * A session whose page statement answers from `outcomes` in order — a row, or an error to throw —
+ * and records the cursor each page was bound to. `beforePage` runs with the page's index before
+ * it answers.
+ */
+function sessionOf(outcomes: Array<PageRow | Error>, beforePage?: (index: number) => void) {
+  const cursors: string[] = []
+  const tx = {
+    unsafe: vi.fn(async (query: string, params?: unknown[]) => {
+      if (!query.includes('WITH page')) return []
+      beforePage?.(cursors.length)
+      cursors.push(String(params?.[0]))
+      const outcome = outcomes.shift()
+      if (outcome === undefined) throw new Error('No outcome left for this page')
+      if (outcome instanceof Error) throw outcome
+      return [outcome]
+    }),
+  }
+  const session = {
+    begin: vi.fn(async (work: (tx: unknown) => Promise<unknown>) => work(tx)),
+    unsafe: vi.fn(async () => []),
+  } as unknown as Sql
+  return { session, cursors }
+}
+
+/** Runs a backfill under fake timers, so its retry pauses pass without waiting. */
+async function backfillNow(...args: Parameters<typeof backfillProjectionSourceAcl>) {
+  vi.useFakeTimers()
+  const result = backfillProjectionSourceAcl(...args)
+  /** A rejection must not surface as unhandled while the timers are still being drained. */
+  const settled = result.catch(() => undefined)
+  await vi.runAllTimersAsync()
+  await settled
+  return result
+}
+
 describe('backfillProjectionSourceAcl', () => {
+  afterEach(() => {
+    vi.useRealTimers()
+  })
+
+  it.each(['55P03', '57014'])(
+    'retries the page after a %s timeout and moves the cursor only once it commits',
+    async (code) => {
+      const { session, cursors } = sessionOf([
+        { scanned: 2, filled: 2, last_id: 'id-2' },
+        postgresError(code),
+        postgresError(code),
+        { scanned: 1, filled: 1, last_id: 'id-3' },
+        { scanned: 0, filled: 0, last_id: null },
+      ])
+      await expect(
+        backfillNow(session, 'embedding_keyword_tin', { pauseMs: 0 })
+      ).resolves.toMatchObject({ scanned: 3, written: 3, afterId: 'id-3', done: true })
+      expect(cursors).toEqual(['', 'id-2', 'id-2', 'id-2', 'id-3'])
+    }
+  )
+
+  it('gives up on a page that times out more than the retry limit in a row', async () => {
+    const { session, cursors } = sessionOf(
+      Array.from({ length: PROJECTION_SOURCE_ACL_PAGE_RETRIES + 1 }, () => postgresError('55P03'))
+    )
+    await expect(backfillNow(session, 'embedding_keyword_tin', { pauseMs: 0 })).rejects.toThrow(
+      'SQLSTATE 55P03'
+    )
+    expect(cursors).toHaveLength(PROJECTION_SOURCE_ACL_PAGE_RETRIES + 1)
+    expect(new Set(cursors)).toEqual(new Set(['']))
+  })
+
+  it('propagates an error that is not a timeout without retrying', async () => {
+    const { session, cursors } = sessionOf([postgresError('42P01')])
+    await expect(backfillNow(session, 'embedding_search', { pauseMs: 0 })).rejects.toThrow(
+      'SQLSTATE 42P01'
+    )
+    expect(cursors).toEqual([''])
+  })
+
+  it('propagates an explicit cancellation, which shares the statement timeout SQLSTATE', async () => {
+    const { session, cursors } = sessionOf([
+      postgresError('57014', 'canceling statement due to user request'),
+    ])
+    await expect(backfillNow(session, 'embedding_search', { pauseMs: 0 })).rejects.toThrow(
+      'user request'
+    )
+    expect(cursors).toEqual([''])
+  })
+
+  it('does not start another page when the budget ran out during the retry pause', async () => {
+    const { session, cursors } = sessionOf([
+      { scanned: 1, filled: 1, last_id: 'id-1' },
+      postgresError('57014'),
+    ])
+    await expect(
+      backfillNow(session, 'embedding_search', { pauseMs: 0, budgetMs: 1000 })
+    ).resolves.toMatchObject({ afterId: 'id-1', done: false })
+    expect(cursors).toEqual(['', 'id-1'])
+  })
+
+  it('leaves a page still failing at the budget to the continuation, from the last committed page', async () => {
+    const { session, cursors } = sessionOf(
+      [{ scanned: 1, filled: 1, last_id: 'id-1' }, postgresError('57014')],
+      /** The second page spends the budget before the database cancels it. */
+      (index) => {
+        if (index === 1) vi.advanceTimersByTime(1000)
+      }
+    )
+    await expect(
+      backfillNow(session, 'embedding_search', { pauseMs: 0, budgetMs: 1000 })
+    ).resolves.toMatchObject({ afterId: 'id-1', done: false })
+    expect(cursors).toEqual(['', 'id-1'])
+  })
+
   it.each([0, -1, 1.5, Number.NaN, Number.POSITIVE_INFINITY])(
     'refuses a page size of %s instead of reporting the projection filled',
     async (pageSize) => {