improvement(knowledge): match keyword chunks before authorizing them

apps/sim/lib/knowledge/__integration__/search-latency.integration.ts

@@ -477,7 +477,7 @@ async function sample(
         item.query.includes('CROSS JOIN LATERAL') ||
         isVectorCandidateQuery(item.query) ||
         item.query.includes('WITH scored_search_candidates') ||
-        item.query.includes('WITH visible_keyword_documents'))
+        item.query.includes('WITH matched_keyword_chunks'))
   )
   const plans: Array<
     CapturedQuery & {
@@ -550,7 +550,7 @@ async function sample(
         assertIndexedCandidates(parsedPlan[0].Plan, diagnostics.vectorCandidateLimit!, width)
       }
     }
-    if (query.query.includes('WITH visible_keyword_documents')) {
+    if (query.query.includes('WITH matched_keyword_chunks')) {
       assertScalarKeywordSorts(parsedPlan[0].Plan)
     }
   }

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

@@ -884,7 +884,7 @@ describe('live repository authorization follows ranked candidates', () => {
       const statement = render(query).sql
       if (statement.includes('AS visible')) return candidatePages.shift() ?? []
       if (statement.includes('WITH scored_search_candidates')) return rerankPages.shift() ?? []
-      if (statement.includes('WITH visible_keyword_documents')) return keywordPages.shift() ?? []
+      if (statement.includes('WITH matched_keyword_chunks')) return keywordPages.shift() ?? []
       if (isExactRanking(statement)) return exactPages.shift() ?? []
       if (statement.includes('AS id FROM')) return probePages.shift() ?? []
       return []
@@ -1102,8 +1102,8 @@ describe('live repository authorization follows ranked candidates', () => {
       expect(getForConnectors).toHaveBeenCalledWith(['allowed-source'], undefined)
       if (mode === 'keyword') {
         const ranking = render(dbChainMockFns.execute.mock.calls[0][0]).sql
-        expect(ranking).toContain('scored_keyword_candidates AS MATERIALIZED')
-        expect(ranking).toContain('ORDER BY keyword_rank DESC, id LIMIT')
+        expect(ranking).toContain('matched_keyword_chunks AS MATERIALIZED')
+        expect(ranking).toContain('ORDER BY keyword_rank DESC, matched_keyword_chunks.id')
         expect(ranking).not.toContain('<=>')
         expect(ranking).not.toContain('"content"')
       } else if (mode === 'tags') {
@@ -1240,6 +1240,22 @@ describe('live repository authorization follows ranked candidates', () => {
     expect(refillPredicate).toContain('revoked-source')
   })
 
+  it('matches keyword chunks before the visibility predicate and ranks only what survives it', async () => {
+    keywordPages.push([candidate('selected', 'allowed-source')])
+    queueTableRows(schemaMock.embedding, [{ id: 'selected', content: 'verified result' }])
+    await executeKeywordSearch({ ...params, query: 'release', queryVector: params.queryVector! })
+    const ranking = render(dbChainMockFns.execute.mock.calls[0][0]).sql
+    const matched = ranking.indexOf('matched_keyword_chunks AS MATERIALIZED')
+    const visible = ranking.indexOf('visible_keyword_documents AS MATERIALIZED')
+    expect(matched).toBeGreaterThanOrEqual(0)
+    expect(visible).toBeGreaterThan(matched)
+    expect(ranking.slice(matched, visible)).not.toContain('keyword_rank')
+    expect(ranking.slice(visible)).toContain('FROM matched_keyword_chunks INNER JOIN')
+    /** The predicate fragments are parameterized, so the restriction is read off the query tree. */
+    const fragments = JSON.stringify(dbChainMockFns.execute.mock.calls[0][0])
+    expect(fragments).toContain('= ANY (ARRAY(SELECT document_id FROM matched_keyword_chunks))')
+  })
+
   it('recomputes keyword candidates after excluding a revoked source and rechecks content access', async () => {
     getForConnectors.mockResolvedValueOnce(identity)
     keywordPages.push(

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

@@ -1100,6 +1100,19 @@ export interface KeywordSearchParams {
  * 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.
+ *
+ * The live-scope ranking query runs in three stages: match, authorize, rank. The
+ * visibility predicate carries correlated subqueries — one per connector, one per
+ * search-integration decision — so evaluating it across a base ahead of the query costs a table
+ * pass priced by how many documents the base holds rather than by how many the query matched.
+ * Matching first restricts that predicate to the documents the query actually matched.
+ *
+ * Two details keep that ordering from paying the saving back. Restricting the predicate with
+ * `document.id = ANY (...)` rather than a subquery keeps the narrowed lookup on a bitmap scan,
+ * which prefetches, where a plain `IN (SELECT ...)` plans as an index walk that does not. And
+ * the match stage carries identifiers only: ranking every match rather than every *visible*
+ * match would detoast one text-search vector per match, which on a mid-frequency term costs
+ * more than the pass it replaces.
  */
 export async function executeKeywordSearch(params: KeywordSearchParams): Promise<SearchResult[]> {
   const { knowledgeBaseIds, topK, query, queryVector, structuredFilters, access } = params
@@ -1133,38 +1146,41 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise
       selectPage: async (limit, offset, excludedSources) => {
         const candidates = await runSearchQuery(params.budget, 'keyword.sql', (executor) =>
           executor.execute<SearchReadCandidate>(sql`
-            WITH visible_keyword_documents AS MATERIALIZED (
-              SELECT ${document.id} AS id FROM ${document}
-              WHERE ${and(
-                inArray(document.knowledgeBaseId, knowledgeBaseIds),
-                ...getDocumentVisibilityConditions(
-                  access,
-                  params.filters,
-                  knowledgeMetadataCandidateAccessCondition(access)
-                ),
-                excludeSearchSources(excludedSources)
-              )}
-            ), scored_keyword_candidates AS MATERIALIZED (
+            WITH matched_keyword_chunks AS MATERIALIZED (
               SELECT ${embeddingKeywordSearch.id} AS id,
-                ${embeddingKeywordSearch.documentId} AS document_id,
-                ${candidateRank} AS keyword_rank
+                ${embeddingKeywordSearch.documentId} AS document_id
               FROM ${embeddingKeywordSearch}
               WHERE ${and(
                 inArray(embeddingKeywordSearch.knowledgeBaseId, knowledgeBaseIds),
                 eq(embeddingKeywordSearch.enabled, true),
                 sql`${embeddingKeywordSearch.contentTsv} @@ ${tsQuery}`,
-                sql`${embeddingKeywordSearch.documentId} IN (SELECT id FROM visible_keyword_documents)`,
-                sql`EXISTS (SELECT 1 FROM visible_keyword_documents)`,
                 tagFilterConditions.length
                   ? sql`EXISTS (
                   SELECT 1 FROM ${embedding} WHERE ${embedding.id} = ${embeddingKeywordSearch.id}
                     AND ${and(...tagFilterConditions)}
                 )`
                   : undefined
               )}
+            ), visible_keyword_documents AS MATERIALIZED (
+              SELECT ${document.id} AS id FROM ${document}
+              WHERE ${and(
+                inArray(document.knowledgeBaseId, knowledgeBaseIds),
+                sql`${document.id} = ANY (ARRAY(SELECT document_id FROM matched_keyword_chunks))`,
+                ...getDocumentVisibilityConditions(
+                  access,
+                  params.filters,
+                  knowledgeMetadataCandidateAccessCondition(access)
+                ),
+                excludeSearchSources(excludedSources)
+              )}
             ), ranked_keyword_candidates AS MATERIALIZED (
-              SELECT * FROM scored_keyword_candidates
-              ORDER BY keyword_rank DESC, id LIMIT ${limit} OFFSET ${offset}
+              SELECT matched_keyword_chunks.id, matched_keyword_chunks.document_id,
+                ${candidateRank} AS keyword_rank
+              FROM matched_keyword_chunks INNER JOIN ${embeddingKeywordSearch}
+                ON ${embeddingKeywordSearch.id} = matched_keyword_chunks.id
+              WHERE matched_keyword_chunks.document_id IN (SELECT id FROM visible_keyword_documents)
+              ORDER BY keyword_rank DESC, matched_keyword_chunks.id
+              LIMIT ${limit} OFFSET ${offset}
             )
             SELECT ranked_keyword_candidates.id, ${document.id} AS "documentId",
               ${document.connectorId} AS "connectorId",