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",