improvement(knowledge): index the projection by source; rank a narrow reader's keyword window once

- a partial btree on embedding_search.connector_id lets the planner lead with a few sources where the caller's tokens alone would match most of the index, which is what made a small slice cost a bitmap over all of it
- a narrow reader's keyword page returns several pages' worth of readable rows from one wide ranking, instead of ranking the window again for each page

Author: waleedlatif1

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

@@ -1574,10 +1574,11 @@ describe('permitted-document planner', () => {
       expect(statement).toContain('ranked_tin_chunks.acl')
     })
 
-    it('takes one wide window for a resolved scope and leaves a short page short', async () => {
+    it('takes one wide window for a narrow resolved scope and leaves a short page short', async () => {
       tinPages = [{ ranked: 50_000, candidates: [hit('a', 'src-a')] }]
       queueTableRows(schemaMock.embedding, [{ ...hit('a', 'src-a'), content: 'release notes' }])
       const results = await keyword({
+        permitted: { kind: 'unbounded', broad: false },
         accessPlan: {
           connectors: { workspace: [], admin: ['src-a'], members: [] },
           observers: { confirmed: [], observed: [] },
@@ -1589,6 +1590,8 @@ describe('permitted-document planner', () => {
       expect(tinStatements()).toHaveLength(1)
       expect(JSON.stringify(tinStatements()[0])).toContain('50000')
       expect(ginStatements()).toHaveLength(0)
+      /** The window is ranked once for several pages' worth of readable rows, not once per page. */
+      expect(JSON.stringify(tinStatements()[0])).toContain('1000')
     })
 
     it('widens the ranked window while too few ranked chunks are readable', async () => {

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

@@ -516,6 +516,9 @@ const FTS_CONFIG = 'english'
  */
 const TIN_KEYWORD_WINDOWS = [2000, 10_000, 50_000] as const
 
+/** Readable rows one wide window returns for a narrow reader: several pages' worth, ranked once. */
+const NARROW_KEYWORD_PAGE = 1000
+
 /**
  * Row visibility predicates shared by every search leg: a chunk is only
  * retrievable when both it and its document are enabled, the document finished
@@ -1837,10 +1840,17 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise
        * one, so that is the only window tried. Either way what a resolved scope cannot fill is
        * left short rather than handed to a ranking over every match.
        */
-      const windows =
-        accessPlan && params.permitted?.kind === 'unbounded' && !params.permitted.broad
-          ? [TIN_KEYWORD_WINDOWS.at(-1)!]
-          : TIN_KEYWORD_WINDOWS
+      const narrow =
+        accessPlan !== undefined &&
+        params.permitted?.kind === 'unbounded' &&
+        !params.permitted.broad
+      const windows = narrow ? [TIN_KEYWORD_WINDOWS.at(-1)!] : TIN_KEYWORD_WINDOWS
+      /**
+       * A narrow reader's page is the readable remainder of a wide ranking, and that ranking is
+       * the cost: each page would rank the window again to find the next few readable rows, so one
+       * statement returns as many as several pages could ask for.
+       */
+      const pageLimit = narrow ? Math.max(limit, NARROW_KEYWORD_PAGE) : limit
       for (const window of windows) {
         if (window < offset + limit) continue
         const [page] = await runSearchQuery(params.budget, 'keyword.tin', (executor) =>
@@ -1868,7 +1878,7 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise
               FROM ranked_tin_chunks /* on-row visibility */
               WHERE ranked_tin_chunks.enabled AND ${onRowKeywordVisibility(excludedSources)}
               ORDER BY ranked_tin_chunks.keyword_rank DESC, ranked_tin_chunks.id
-              LIMIT ${limit} OFFSET ${offset}`
+              LIMIT ${pageLimit} OFFSET ${offset}`
                   : sql`
               SELECT ranked_tin_chunks.id, ${document.id} AS "documentId",
                 ${document.connectorId} AS "connectorId",
@@ -1878,7 +1888,7 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise
               WHERE ranked_tin_chunks.enabled
                 AND ${and(...candidateDocumentConditions(knowledgeBaseIds, access, params.filters, knowledgeAccessCondition(access)), excludeSearchSources(excludedSources))}
               ORDER BY ranked_tin_chunks.keyword_rank DESC, ranked_tin_chunks.id
-              LIMIT ${limit} OFFSET ${offset}`
+              LIMIT ${pageLimit} OFFSET ${offset}`
               }
             )
             SELECT (SELECT count(*)::int FROM ranked_tin_chunks) AS ranked,

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

@@ -129,9 +129,11 @@ export async function backfillProjectionSourceAcl(
 }
 
 /**
- * The ACL index each projection needs for exact ranking of a small readable set, built after the bulk
- * load and concurrently, so the triggers keep writing. `CONCURRENTLY` cannot run in a transaction, and
- * the pool's lock timeout would cancel a build that merely waits for a long transaction to finish.
+ * The indexes exact ranking of a readable set needs: the ACL index on each projection, and on the
+ * vector projection the source index that lets the planner lead with a few sources when the
+ * caller's tokens alone would match most of the index. Built after the bulk load and concurrently,
+ * so the triggers keep writing. `CONCURRENTLY` cannot run in a transaction, and the pool's lock
+ * timeout would cancel a build that merely waits for a long transaction to finish.
  */
 export async function indexProjectionAcl(sql: Sql): Promise<void> {
   const [{ timeout }] = await sql`SELECT current_setting('lock_timeout') AS timeout`
@@ -143,6 +145,10 @@ export async function indexProjectionAcl(sql: Sql): Promise<void> {
           ON ${projection} USING gin (acl) WHERE enabled`
       )
     }
+    await sql.unsafe(
+      `CREATE INDEX CONCURRENTLY IF NOT EXISTS embedding_search_source_idx
+        ON embedding_search (connector_id) WHERE enabled`
+    )
   } finally {
     await sql`SELECT set_config('lock_timeout', ${timeout}, false)`
   }