fix(knowledge): walk a large bounded set on the row before ranking it exactly

A member reading most of a large source, with that source selected as a filter, enumerated a
bounded permitted set of tens of thousands of documents and then ranked every chunk of it
exactly on both legs: the vector leg read every chunk's projected vector, and the keyword leg
materialized every chunk of the set before it matched the term. Cold, each leg outran its budget
and the search returned nothing.

A bounded set past a size limit is now walked on the row first, where the plan's source and ACL
decide readability and the walk stops at its tuple cap, and ranked exactly only when the walk
cannot fill its pool, so recall is never below the exact ranking's. The keyword leg treats the
same set as a narrow on-row reader: Tin windows where Tin serves, otherwise the GIN shape whose
cost follows the term's matches. Sets under the limit keep their exact paths.

Author: waleedlatif1

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

@@ -41,6 +41,7 @@ import {
   handleTagAndVectorSearch,
   handleTagOnlySearch,
   handleVectorOnlySearch,
+  PERMITTED_EXACT_DOCUMENT_LIMIT,
   type PermittedDocuments,
   resolvePermittedDocuments,
   resolveReach,
@@ -1794,6 +1795,51 @@ describe('permitted-document planner', () => {
       expect(ginStatements()).toHaveLength(0)
     })
 
+    describe('a bounded set past the exact-ranking size', () => {
+      const large = Array.from({ length: PERMITTED_EXACT_DOCUMENT_LIMIT }, (_, index) => ({
+        id: `doc-${index}`,
+        connectorId: 'src-a',
+      }))
+      const accessPlan = {
+        connectors: { workspace: [], admin: ['src-a'], members: [], liveProofRequired: [] },
+        observers: { confirmed: [], observed: [] },
+        memberSources: [],
+        connectorTypes: new Map(),
+        uploads: true,
+      }
+
+      it('ranks with Tin as a narrow reader, decided on the row', async () => {
+        tinPages = [{ ranked: 1500, candidates: [hit('a', 'src-a')] }]
+        queueTableRows(schemaMock.embedding, [{ ...hit('a', 'src-a'), content: 'release notes' }])
+        const results = await keyword({
+          permitted: { kind: 'bounded', documents: large },
+          accessPlan,
+        })
+        expect(results.map((row) => row.id)).toEqual(['a'])
+        expect(mockResolveTinKeywordQuery).toHaveBeenCalledTimes(1)
+        expect(tinStatements()).toHaveLength(1)
+        expect(JSON.stringify(tinStatements()[0])).toContain('2000')
+        expect(JSON.stringify(tinStatements()[0])).not.toContain('doc-4999')
+        expect(ginStatements()).toHaveLength(0)
+      })
+
+      it('falls back to a GIN ranking that reads what the term matches, not every chunk of the set', async () => {
+        mockResolveTinKeywordQuery.mockResolvedValue(null)
+        queueTableRows(schemaMock.embedding, [{ ...hit('a', 'src-a'), content: 'release notes' }])
+        await keyword({ permitted: { kind: 'bounded', documents: large }, accessPlan })
+        expect(tinStatements()).toHaveLength(0)
+        expect(ginStatements()).toHaveLength(1)
+        expect(JSON.stringify(ginStatements()[0])).not.toContain('doc-4999')
+      })
+
+      it('keeps the bounded read for a set under the size', async () => {
+        mockResolveTinKeywordQuery.mockResolvedValue(null)
+        await keyword({ permitted: { kind: 'bounded', documents: large.slice(0, -1) }, accessPlan })
+        expect(mockResolveTinKeywordQuery).not.toHaveBeenCalled()
+        expect(JSON.stringify(ginStatements()[0])).toContain('doc-4998')
+      })
+    })
+
     it('leaves the page to the GIN ranking when the widest window cannot fill it', async () => {
       tinPages = [
         { ranked: 2000, candidates: [] },
@@ -2400,6 +2446,59 @@ describe('filters on a resolved scope', () => {
     expect(JSON.stringify(walks[0])).toContain('release')
   })
 
+  describe('a bounded set past the exact-ranking size', () => {
+    const large = Array.from({ length: PERMITTED_EXACT_DOCUMENT_LIMIT }, (_, index) => ({
+      id: `doc-${index}`,
+      connectorId: 'src-a',
+    }))
+    const walked = Array.from({ length: 200 }, (_, index) => hit(`w-${index}`, 'src-a'))
+    const search = (documents: typeof large) =>
+      handleVectorOnlySearch({
+        ...params,
+        permitted: { kind: 'bounded', documents },
+        accessPlan: plan(),
+      })
+    beforeEach(() => {
+      const execute = dbChainMockFns.execute.getMockImplementation()!
+      dbChainMockFns.execute.mockImplementation(async (query) => {
+        /** The projection is filled, so a walk decides readability on the row. */
+        if (render(query).sql.includes('AS unfilled')) return [{ unfilled: false }]
+        return execute(query)
+      })
+    })
+
+    it('walks the graph on the row instead of ranking every chunk of the set', async () => {
+      traversedRows = walked
+      queueTableRows(schemaMock.embedding, [walked[0]])
+      expect((await search(large)).map((row) => row.id)).toEqual(['w-0'])
+      const walks = statements().filter((query) => isWalk(query.sql))
+      expect(walks).toHaveLength(1)
+      expect(statements().filter((query) => isExactRanking(query.sql))).toHaveLength(0)
+      /** Readability rides on the row through the plan; the set's identifiers never cross the wire. */
+      expect(JSON.stringify(walks[0])).not.toContain('doc-4999')
+      expect(JSON.stringify(walks[0])).toContain('src-a')
+    })
+
+    it('ranks the set exactly when the walk cannot fill its pool', async () => {
+      traversedRows = []
+      exactRows = [hit('a', 'src-a')]
+      queueTableRows(schemaMock.embedding, [hit('a', 'src-a')])
+      expect((await search(large)).map((row) => row.id)).toEqual(['a'])
+      expect(statements().filter((query) => isWalk(query.sql))).toHaveLength(1)
+      const exact = statements().filter((query) => isExactRanking(query.sql))
+      expect(exact).toHaveLength(1)
+      expect(JSON.stringify(exact[0])).toContain('doc-4999')
+    })
+
+    it('ranks a set under the size exactly, without a walk', async () => {
+      exactRows = [hit('a', 'src-a')]
+      queueTableRows(schemaMock.embedding, [hit('a', 'src-a')])
+      expect((await search(large.slice(0, -1))).map((row) => row.id)).toEqual(['a'])
+      expect(statements().filter((query) => isWalk(query.sql))).toHaveLength(0)
+      expect(statements().filter((query) => isExactRanking(query.sql))).toHaveLength(1)
+    })
+  })
+
   it('ranks a date-bounded set exactly even when a member source has its own index', async () => {
     indexedSourceRows = [{ name: 'idx', connectorId: 'member-src' }]
     exactRows = [{ id: 'a' }]

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

@@ -216,6 +216,16 @@ export const VECTOR_PROBE_DOCUMENT_LIMIT = Math.round(
   (VECTOR_PROBE_BUDGET_MS * 1000) / VECTOR_PROBE_MICROSECONDS_PER_DOCUMENT
 )
 
+/**
+ * Documents a bounded permitted set may hold before ranking it exactly costs more than walking
+ * the graph on the row. Exact ranking reads every chunk of the set, a few per document, where an
+ * on-row walk reads at most {@link CANDIDATE_HNSW_MAX_SCAN_TUPLES} tuples; at this size the two
+ * meet. A set past it is walked first and ranked exactly only if the walk cannot fill its pool, so
+ * its recall is never below the exact ranking's and its usual cost is the walk's. The same size
+ * turns the keyword leg from a read of the set's every chunk into a ranking decided on the row.
+ */
+export const PERMITTED_EXACT_DOCUMENT_LIMIT = 5_000
+
 /** How long to stop trying the iterative-scan settings after the server rejected them. */
 const HNSW_SETTINGS_UNSUPPORTED_RETRY_MS = 10 * 60 * 1000
 
@@ -1766,6 +1776,49 @@ async function selectVectorResults(params: SearchParams): Promise<SearchResult[]
           )
         }
         let selected: SearchReadCandidate[]
+        /**
+         * The bounded ANN traversal is the whole candidate set. LIMIT keeps document
+         * authorization downstream of the traversal, with a primary-key lookup per candidate.
+         */
+        const scopeOfWalk = and(
+          inArray(embeddingSearch.knowledgeBaseId, params.knowledgeBaseIds),
+          eq(embeddingSearch.enabled, true),
+          excludedOnRow
+        )
+        const walkGraph = () =>
+          withVectorScanSettings(
+            (executor) =>
+              executor.execute<SearchReadCandidate>(
+                plan
+                  ? sql`
+            SELECT ${PROJECTION_CANDIDATE_COLUMNS}
+            FROM ${embeddingSearch} /* on-row visibility */
+            WHERE ${and(
+              scopeOfWalk,
+              projectionCandidateAccessCondition(embeddingSearch, params.access, plan, { filled }),
+              documentCondition === undefined
+                ? undefined
+                : sql`EXISTS (SELECT 1 FROM ${document} WHERE ${and(eq(document.id, embeddingSearch.documentId), documentCondition)})`
+            )}
+            ORDER BY ${candidateDistance} LIMIT ${candidateLimit}
+          `
+                  : sql`
+            SELECT ${embeddingSearch.id} AS id, ${embeddingSearch.documentId} AS "documentId",
+              visible.connector_id AS "connectorId"
+            FROM ${embeddingSearch}
+            CROSS JOIN LATERAL (
+              SELECT ${document.connectorId} AS connector_id FROM ${document}
+              WHERE ${and(eq(document.id, embeddingSearch.documentId), ...candidateDocumentVisibility, candidateTagCondition)}
+              LIMIT 1
+            ) AS visible
+            WHERE ${scopeOfWalk}
+            ORDER BY ${candidateDistance} LIMIT ${candidateLimit}
+          `
+              ),
+            params.budget,
+            'vector.candidate_search',
+            plan ? onRowWalkScanTuples(documentCondition, filled) : undefined
+          )
         /**
          * A source the caller is a member of that has its own index is walked on its own, which
          * beats ranking it exactly once it is large enough to have earned that index.
@@ -1774,6 +1827,25 @@ async function selectVectorResults(params: SearchParams): Promise<SearchResult[]
           (id) => plannedIndexedSources?.has(id) ?? false
         )
         if (
+          params.permitted?.kind === 'bounded' &&
+          plan &&
+          filled &&
+          params.permitted.documents.length >= PERMITTED_EXACT_DOCUMENT_LIMIT
+        ) {
+          /**
+           * A set this large costs more to rank exactly than to walk: exact ranking reads every
+           * chunk of every document in it, while the walk decides readability on the rows it
+           * visits and stops at its tuple cap. The walk answers whenever the set is a fair share
+           * of the graph; where it is not, the walk underfills and the exact ranking that was
+           * always complete takes over, so nothing is lost but the walk's bounded cost.
+           */
+          selected = await walkGraph()
+          if (selected.length < candidateLimit) {
+            selected = await rankPermittedExactly(
+              params.permitted.documents.map((entry) => entry.id)
+            )
+          }
+        } else if (
           params.permitted?.kind === 'bounded' &&
           (!walksASource || dateFilterCondition(params.filters) || params.filters?.source)
         ) {
@@ -1806,48 +1878,7 @@ async function selectVectorResults(params: SearchParams): Promise<SearchResult[]
             budget: params.budget,
           })
         } else {
-          /**
-           * The bounded ANN traversal is the whole candidate set. LIMIT keeps document
-           * authorization downstream of the traversal, with a primary-key lookup per candidate.
-           */
-          const scopeOfWalk = and(
-            inArray(embeddingSearch.knowledgeBaseId, params.knowledgeBaseIds),
-            eq(embeddingSearch.enabled, true),
-            excludedOnRow
-          )
-          selected = await withVectorScanSettings(
-            (executor) =>
-              executor.execute<SearchReadCandidate>(
-                plan
-                  ? sql`
-            SELECT ${PROJECTION_CANDIDATE_COLUMNS}
-            FROM ${embeddingSearch} /* on-row visibility */
-            WHERE ${and(
-              scopeOfWalk,
-              projectionCandidateAccessCondition(embeddingSearch, params.access, plan, { filled }),
-              documentCondition === undefined
-                ? undefined
-                : sql`EXISTS (SELECT 1 FROM ${document} WHERE ${and(eq(document.id, embeddingSearch.documentId), documentCondition)})`
-            )}
-            ORDER BY ${candidateDistance} LIMIT ${candidateLimit}
-          `
-                  : sql`
-            SELECT ${embeddingSearch.id} AS id, ${embeddingSearch.documentId} AS "documentId",
-              visible.connector_id AS "connectorId"
-            FROM ${embeddingSearch}
-            CROSS JOIN LATERAL (
-              SELECT ${document.connectorId} AS connector_id FROM ${document}
-              WHERE ${and(eq(document.id, embeddingSearch.documentId), ...candidateDocumentVisibility, candidateTagCondition)}
-              LIMIT 1
-            ) AS visible
-            WHERE ${scopeOfWalk}
-            ORDER BY ${candidateDistance} LIMIT ${candidateLimit}
-          `
-              ),
-            params.budget,
-            'vector.candidate_search',
-            plan ? onRowWalkScanTuples(documentCondition, filled) : undefined
-          )
+          selected = await walkGraph()
           /**
            * A full traversal is already the nearest permitted chunks, so nothing else is worth
            * running. An underfilled one is the signal that visibility removed neighbours the graph
@@ -2021,10 +2052,18 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise
      * A caller reaching past the permitted-set limit reads much of the index, so ranking every
      * match before checking access is the leg's whole cost for a common term. Where the Tin
      * projection is complete, BM25 ranks inside the bases first and access is checked only on the
-     * top of that ranking.
+     * top of that ranking. A bounded set past the exact-ranking size is read on the row like an
+     * unbounded one: the bounded read materializes every chunk of the set before it matches a
+     * term, where a ranking decided on the row costs what the term matches.
      */
+    const accessPlan = access.kind === 'user' ? params.accessPlan : undefined
+    const largePermittedSet =
+      accessPlan !== undefined &&
+      params.permitted?.kind === 'bounded' &&
+      params.permitted.documents.length >= PERMITTED_EXACT_DOCUMENT_LIMIT
+    const onRowReader = params.permitted?.kind === 'unbounded' || largePermittedSet
     let tinQuery: Awaited<ReturnType<typeof resolveTinKeywordQuery>> = null
-    if (params.permitted?.kind === 'unbounded' && tagFilterConditions.length === 0) {
+    if (onRowReader && tagFilterConditions.length === 0) {
       try {
         tinQuery = await resolveTinKeywordQuery(
           params.searchIndexOnly === true,
@@ -2038,9 +2077,7 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise
         return []
       }
     }
-    if (params.permitted?.kind === 'unbounded')
-      annotateSearchDiagnostics({ keywordRanking: tinQuery ? 'tin' : 'gin' })
-    const accessPlan = access.kind === 'user' ? params.accessPlan : undefined
+    if (onRowReader) annotateSearchDiagnostics({ keywordRanking: tinQuery ? 'tin' : 'gin' })
     /** A filled projection decides readability on the ranked row alone; none of its rows needs the document. */
     const tinFilled =
       accessPlan && tinQuery
@@ -2098,8 +2135,7 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise
        */
       const narrow =
         accessPlan !== undefined &&
-        params.permitted?.kind === 'unbounded' &&
-        !params.permitted.broad
+        ((params.permitted?.kind === 'unbounded' && !params.permitted.broad) || largePermittedSet)
       const windows: readonly number[] = narrow ? NARROW_KEYWORD_WINDOWS : TIN_KEYWORD_WINDOWS
       /**
        * A narrow reader's page is the readable remainder of a wide ranking, and that ranking is
@@ -2190,7 +2226,7 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise
          * still re-applies the candidate predicate, so the restriction can only narrow.
          */
         const permittedIds =
-          params.permitted?.kind === 'bounded'
+          params.permitted?.kind === 'bounded' && !largePermittedSet
             ? params.permitted.documents.map((entry) => entry.id)
             : undefined
         if (permittedIds?.length === 0) return { candidates: [], nextOffset: offset }