From 9f0e816a816b7dee43ad46b2fd22588cb64d5a18 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Fri, 18 Sep 2026 11:05:57 -0700 Subject: [PATCH] improvement(knowledge): match keyword chunks before authorizing them The org-scoped keyword leg built its candidate set in the order visibility -> match -> rank. The visibility predicate carries a correlated subquery per connector plus a search-integration check, so evaluating it across the base before the query terms were consulted priced every search by how many documents the base holds rather than by how many the query matched. Reorder to match -> authorize -> rank. The match stage carries chunk identifiers only, the identical visibility predicate then runs over just the documents that matched, and ts_rank_cd is computed for the matches that survive it. Same predicate, same ordering, same page contract. Two details keep the reorder 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 ranking in the match stage rather than after authorization would detoast one text-search vector per match, which on a mid-frequency term costs more than the pass it replaces. --- .../search-latency.integration.ts | 4 +- apps/sim/lib/knowledge/search/queries.test.ts | 22 ++++++-- apps/sim/lib/knowledge/search/queries.ts | 52 ++++++++++++------- 3 files changed, 55 insertions(+), 23 deletions(-) diff --git a/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts b/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts index e555d9af3e7..68981209d69 100644 --- a/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts +++ b/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) } } diff --git a/apps/sim/lib/knowledge/search/queries.test.ts b/apps/sim/lib/knowledge/search/queries.test.ts index ce1f0031c63..984829a43dc 100644 --- a/apps/sim/lib/knowledge/search/queries.test.ts +++ b/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( diff --git a/apps/sim/lib/knowledge/search/queries.ts b/apps/sim/lib/knowledge/search/queries.ts index bb5ae17f674..1792f735f72 100644 --- a/apps/sim/lib/knowledge/search/queries.ts +++ b/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 { const { knowledgeBaseIds, topK, query, queryVector, structuredFilters, access } = params @@ -1133,28 +1146,14 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise selectPage: async (limit, offset, excludedSources) => { const candidates = await runSearchQuery(params.budget, 'keyword.sql', (executor) => executor.execute(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} @@ -1162,9 +1161,26 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise )` : 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",