perf(knowledge): stop the keyword leg detoasting every match's vector

Selecting the cosine distance in the ranking query made Postgres detoast the
1536-dimension embedding and compute a distance for every full-text match
before the LIMIT applied, so cost tracked how common the query term was rather
than topK. On a 20k-chunk base with a term matching every row that was 61,055
buffer hits against 1,030 for the same query without the projection.

Rank on ids and ts_rank_cd alone, then hydrate only the rows that survive the
limit. Same results, and the worst case drops to ~27ms end to end.

Author: waleedlatif1

apps/sim/app/api/knowledge/search/utils.test.ts

@@ -394,6 +394,31 @@ describe('Knowledge Search Utils', () => {
       expect(dbChainMockFns.select).toHaveBeenCalledTimes(knowledgeBaseIds.length)
     })
 
+    it('ranks without selecting the embedding column, then hydrates the survivors', async () => {
+      queueTableRows(schemaMock.embedding, [{ id: 'kw-1', keywordRank: 0.9 }])
+      queueTableRows(schemaMock.embedding, [makeResult('kw-1')])
+
+      const results = await executeKeywordSearch({
+        knowledgeBaseIds: ['kb-1'],
+        topK: 10,
+        query: 'PROJ-1234',
+        queryVector: JSON.stringify([0.1, 0.2, 0.3]),
+      })
+
+      expect(results.map((r) => r.id)).toEqual(['kw-1'])
+      expect(dbChainMockFns.select).toHaveBeenCalledTimes(2)
+
+      /**
+       * Projecting the distance in the ranking pass makes Postgres detoast the
+       * 1536-dimension vector for every full-text match before the LIMIT, so
+       * cost tracks how common the term is rather than topK. The ranking pass
+       * must select ids and relevance only.
+       */
+      const rankingSelect = dbChainMockFns.select.mock.calls[0][0]
+      expect(Object.keys(rankingSelect)).toEqual(['id', 'keywordRank'])
+      expect(Object.keys(dbChainMockFns.select.mock.calls[1][0])).toContain('distance')
+    })
+
     it('uses a single query when the parallel threshold is not crossed', async () => {
       const knowledgeBaseIds = ['kb-1', 'kb-2']
       expect(getQueryStrategy(knowledgeBaseIds.length, 10).useParallel).toBe(false)
@@ -451,7 +476,9 @@ describe('Knowledge Search Utils', () => {
     })
 
     it('runs both legs and fuses them in hybrid mode', async () => {
+      // Vector leg, then the keyword leg's ranking pass, then its hydration pass.
       queueTableRows(schemaMock.embedding, [makeResult('vector-hit')])
+      queueTableRows(schemaMock.embedding, [{ id: 'keyword-hit', keywordRank: 0.9 }])
       queueTableRows(schemaMock.embedding, [makeResult('keyword-hit')])
 
       const results = await executeKnowledgeSearch({
@@ -463,7 +490,7 @@ describe('Knowledge Search Utils', () => {
       })
 
       expect(results.map((r) => r.id).sort()).toEqual(['keyword-hit', 'vector-hit'])
-      expect(dbChainMockFns.select).toHaveBeenCalledTimes(2)
+      expect(dbChainMockFns.select).toHaveBeenCalledTimes(3)
     })
 
     it('falls back to vector results when the keyword leg fails', async () => {

apps/sim/app/api/knowledge/search/utils.ts

@@ -2,7 +2,7 @@ import { db } from '@sim/db'
 import { document, embedding } from '@sim/db/schema'
 import { createLogger } from '@sim/logger'
 import { getErrorMessage } from '@sim/utils/errors'
-import { and, eq, inArray, isNull, sql } from 'drizzle-orm'
+import { and, eq, inArray, isNull, type SQL, sql } from 'drizzle-orm'
 import type { StructuredFilter } from '@/lib/knowledge/types'
 
 const logger = createLogger('KnowledgeSearch')
@@ -570,6 +570,14 @@ export interface KeywordSearchParams {
  * lexically consume every slot, so an exact-token hit in a smaller base would
  * never reach fusion. Both legs must draw candidates the same way, or rank
  * fusion is combining rankings taken over differently-shaped pools.
+ *
+ * Ranking and hydration are two steps on purpose. Projecting the cosine
+ * distance in the ranking query makes Postgres detoast the 1536-dimension
+ * vector and compute a distance for *every* full-text match before the `LIMIT`
+ * applies — work that scales with how common the query term is rather than
+ * with `topK` (measured at ~59x the buffer reads on a 20k-chunk base for a term
+ * matching every row). Ranking therefore touches no vectors, and only the rows
+ * that survive the limit are hydrated.
  */
 export async function executeKeywordSearch(params: KeywordSearchParams): Promise<SearchResult[]> {
   const { knowledgeBaseIds, topK, query, queryVector, structuredFilters } = params
@@ -584,58 +592,55 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise
     ? getStructuredTagFilters(structuredFilters, embedding)
     : []
 
-  /** Selected alongside the row so per-base batches can be re-ranked globally. */
-  const selectFields = {
-    ...getSearchResultFields(
-      sql<number>`${embedding.embedding} <=> ${queryVector}::vector`.as('distance')
-    ),
-    keywordRank: rankExpr.as('keyword_rank'),
-  }
+  const rankConditions = (kbScope: SQL | undefined) =>
+    and(
+      kbScope,
+      ...getVisibilityConditions(),
+      sql`${embedding.contentTsv} @@ ${tsQuery}`,
+      ...tagFilterConditions
+    )
+
+  /** Ranking pass: ids and relevance only, so no vector is read. */
+  const rankRows = (kbScope: SQL | undefined, limit: number) =>
+    db
+      .select({ id: embedding.id, keywordRank: rankExpr.as('keyword_rank') })
+      .from(embedding)
+      .innerJoin(document, eq(embedding.documentId, document.id))
+      .where(rankConditions(kbScope))
+      .orderBy(sql`${rankExpr} DESC`)
+      .limit(limit)
 
   const strategy = getQueryStrategy(knowledgeBaseIds.length, topK)
 
+  let ranked: { id: string; keywordRank: number }[]
   if (strategy.useParallel) {
     const parallelLimit = Math.ceil(topK / knowledgeBaseIds.length) + 5
-
     const perBase = await Promise.all(
-      knowledgeBaseIds.map((kbId) =>
-        db
-          .select(selectFields)
-          .from(embedding)
-          .innerJoin(document, eq(embedding.documentId, document.id))
-          .where(
-            and(
-              eq(embedding.knowledgeBaseId, kbId),
-              ...getVisibilityConditions(),
-              sql`${embedding.contentTsv} @@ ${tsQuery}`,
-              ...tagFilterConditions
-            )
-          )
-          .orderBy(sql`${rankExpr} DESC`)
-          .limit(parallelLimit)
-      )
+      knowledgeBaseIds.map((kbId) => rankRows(eq(embedding.knowledgeBaseId, kbId), parallelLimit))
     )
+    ranked = perBase.flat().sort((a, b) => b.keywordRank - a.keywordRank)
+  } else {
+    ranked = await rankRows(inArray(embedding.knowledgeBaseId, knowledgeBaseIds), topK)
+  }
 
-    return perBase
-      .flat()
-      .sort((a, b) => b.keywordRank - a.keywordRank)
-      .slice(0, topK)
+  const topIds = ranked.slice(0, topK).map((row) => row.id)
+  if (topIds.length === 0) {
+    return []
   }
 
-  return await db
-    .select(selectFields)
-    .from(embedding)
-    .innerJoin(document, eq(embedding.documentId, document.id))
-    .where(
-      and(
-        inArray(embedding.knowledgeBaseId, knowledgeBaseIds),
-        ...getVisibilityConditions(),
-        sql`${embedding.contentTsv} @@ ${tsQuery}`,
-        ...tagFilterConditions
+  /** Hydration pass: full rows plus the cosine distance, bounded to the survivors. */
+  const hydrated = await db
+    .select(
+      getSearchResultFields(
+        sql<number>`${embedding.embedding} <=> ${queryVector}::vector`.as('distance')
       )
     )
-    .orderBy(sql`${rankExpr} DESC`)
-    .limit(topK)
+    .from(embedding)
+    .innerJoin(document, eq(embedding.documentId, document.id))
+    .where(and(inArray(embedding.id, topIds), ...getVisibilityConditions()))
+
+  const rowById = new Map(hydrated.map((row) => [row.id, row]))
+  return topIds.map((id) => rowById.get(id)).filter((row): row is SearchResult => row !== undefined)
 }
 
 /**