Skip to content

Commit 6904e58

Browse files
waleedlatif1waleedlatif1
authored
improvement(knowledge): match keyword chunks before authorizing them (#7973)
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.
1 parent 5d520ba commit 6904e58

3 files changed

Lines changed: 55 additions & 23 deletions

File tree

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -477,7 +477,7 @@ async function sample(
477477
item.query.includes('CROSS JOIN LATERAL') ||
478478
isVectorCandidateQuery(item.query) ||
479479
item.query.includes('WITH scored_search_candidates') ||
480-
item.query.includes('WITH visible_keyword_documents'))
480+
item.query.includes('WITH matched_keyword_chunks'))
481481
)
482482
const plans: Array<
483483
CapturedQuery & {
@@ -550,7 +550,7 @@ async function sample(
550550
assertIndexedCandidates(parsedPlan[0].Plan, diagnostics.vectorCandidateLimit!, width)
551551
}
552552
}
553-
if (query.query.includes('WITH visible_keyword_documents')) {
553+
if (query.query.includes('WITH matched_keyword_chunks')) {
554554
assertScalarKeywordSorts(parsedPlan[0].Plan)
555555
}
556556
}

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

Lines changed: 19 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -884,7 +884,7 @@ describe('live repository authorization follows ranked candidates', () => {
884884
const statement = render(query).sql
885885
if (statement.includes('AS visible')) return candidatePages.shift() ?? []
886886
if (statement.includes('WITH scored_search_candidates')) return rerankPages.shift() ?? []
887-
if (statement.includes('WITH visible_keyword_documents')) return keywordPages.shift() ?? []
887+
if (statement.includes('WITH matched_keyword_chunks')) return keywordPages.shift() ?? []
888888
if (isExactRanking(statement)) return exactPages.shift() ?? []
889889
if (statement.includes('AS id FROM')) return probePages.shift() ?? []
890890
return []
@@ -1102,8 +1102,8 @@ describe('live repository authorization follows ranked candidates', () => {
11021102
expect(getForConnectors).toHaveBeenCalledWith(['allowed-source'], undefined)
11031103
if (mode === 'keyword') {
11041104
const ranking = render(dbChainMockFns.execute.mock.calls[0][0]).sql
1105-
expect(ranking).toContain('scored_keyword_candidates AS MATERIALIZED')
1106-
expect(ranking).toContain('ORDER BY keyword_rank DESC, id LIMIT')
1105+
expect(ranking).toContain('matched_keyword_chunks AS MATERIALIZED')
1106+
expect(ranking).toContain('ORDER BY keyword_rank DESC, matched_keyword_chunks.id')
11071107
expect(ranking).not.toContain('<=>')
11081108
expect(ranking).not.toContain('"content"')
11091109
} else if (mode === 'tags') {
@@ -1240,6 +1240,22 @@ describe('live repository authorization follows ranked candidates', () => {
12401240
expect(refillPredicate).toContain('revoked-source')
12411241
})
12421242

1243+
it('matches keyword chunks before the visibility predicate and ranks only what survives it', async () => {
1244+
keywordPages.push([candidate('selected', 'allowed-source')])
1245+
queueTableRows(schemaMock.embedding, [{ id: 'selected', content: 'verified result' }])
1246+
await executeKeywordSearch({ ...params, query: 'release', queryVector: params.queryVector! })
1247+
const ranking = render(dbChainMockFns.execute.mock.calls[0][0]).sql
1248+
const matched = ranking.indexOf('matched_keyword_chunks AS MATERIALIZED')
1249+
const visible = ranking.indexOf('visible_keyword_documents AS MATERIALIZED')
1250+
expect(matched).toBeGreaterThanOrEqual(0)
1251+
expect(visible).toBeGreaterThan(matched)
1252+
expect(ranking.slice(matched, visible)).not.toContain('keyword_rank')
1253+
expect(ranking.slice(visible)).toContain('FROM matched_keyword_chunks INNER JOIN')
1254+
/** The predicate fragments are parameterized, so the restriction is read off the query tree. */
1255+
const fragments = JSON.stringify(dbChainMockFns.execute.mock.calls[0][0])
1256+
expect(fragments).toContain('= ANY (ARRAY(SELECT document_id FROM matched_keyword_chunks))')
1257+
})
1258+
12431259
it('recomputes keyword candidates after excluding a revoked source and rechecks content access', async () => {
12441260
getForConnectors.mockResolvedValueOnce(identity)
12451261
keywordPages.push(

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

Lines changed: 34 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1100,6 +1100,19 @@ export interface KeywordSearchParams {
11001100
* with `topK` (measured at ~59x the buffer reads on a 20k-chunk base for a term
11011101
* matching every row). Ranking therefore touches no vectors, and only the rows
11021102
* that survive the limit are hydrated.
1103+
*
1104+
* The live-scope ranking query runs in three stages: match, authorize, rank. The
1105+
* visibility predicate carries correlated subqueries — one per connector, one per
1106+
* search-integration decision — so evaluating it across a base ahead of the query costs a table
1107+
* pass priced by how many documents the base holds rather than by how many the query matched.
1108+
* Matching first restricts that predicate to the documents the query actually matched.
1109+
*
1110+
* Two details keep that ordering from paying the saving back. Restricting the predicate with
1111+
* `document.id = ANY (...)` rather than a subquery keeps the narrowed lookup on a bitmap scan,
1112+
* which prefetches, where a plain `IN (SELECT ...)` plans as an index walk that does not. And
1113+
* the match stage carries identifiers only: ranking every match rather than every *visible*
1114+
* match would detoast one text-search vector per match, which on a mid-frequency term costs
1115+
* more than the pass it replaces.
11031116
*/
11041117
export async function executeKeywordSearch(params: KeywordSearchParams): Promise<SearchResult[]> {
11051118
const { knowledgeBaseIds, topK, query, queryVector, structuredFilters, access } = params
@@ -1133,38 +1146,41 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise
11331146
selectPage: async (limit, offset, excludedSources) => {
11341147
const candidates = await runSearchQuery(params.budget, 'keyword.sql', (executor) =>
11351148
executor.execute<SearchReadCandidate>(sql`
1136-
WITH visible_keyword_documents AS MATERIALIZED (
1137-
SELECT ${document.id} AS id FROM ${document}
1138-
WHERE ${and(
1139-
inArray(document.knowledgeBaseId, knowledgeBaseIds),
1140-
...getDocumentVisibilityConditions(
1141-
access,
1142-
params.filters,
1143-
knowledgeMetadataCandidateAccessCondition(access)
1144-
),
1145-
excludeSearchSources(excludedSources)
1146-
)}
1147-
), scored_keyword_candidates AS MATERIALIZED (
1149+
WITH matched_keyword_chunks AS MATERIALIZED (
11481150
SELECT ${embeddingKeywordSearch.id} AS id,
1149-
${embeddingKeywordSearch.documentId} AS document_id,
1150-
${candidateRank} AS keyword_rank
1151+
${embeddingKeywordSearch.documentId} AS document_id
11511152
FROM ${embeddingKeywordSearch}
11521153
WHERE ${and(
11531154
inArray(embeddingKeywordSearch.knowledgeBaseId, knowledgeBaseIds),
11541155
eq(embeddingKeywordSearch.enabled, true),
11551156
sql`${embeddingKeywordSearch.contentTsv} @@ ${tsQuery}`,
1156-
sql`${embeddingKeywordSearch.documentId} IN (SELECT id FROM visible_keyword_documents)`,
1157-
sql`EXISTS (SELECT 1 FROM visible_keyword_documents)`,
11581157
tagFilterConditions.length
11591158
? sql`EXISTS (
11601159
SELECT 1 FROM ${embedding} WHERE ${embedding.id} = ${embeddingKeywordSearch.id}
11611160
AND ${and(...tagFilterConditions)}
11621161
)`
11631162
: undefined
11641163
)}
1164+
), visible_keyword_documents AS MATERIALIZED (
1165+
SELECT ${document.id} AS id FROM ${document}
1166+
WHERE ${and(
1167+
inArray(document.knowledgeBaseId, knowledgeBaseIds),
1168+
sql`${document.id} = ANY (ARRAY(SELECT document_id FROM matched_keyword_chunks))`,
1169+
...getDocumentVisibilityConditions(
1170+
access,
1171+
params.filters,
1172+
knowledgeMetadataCandidateAccessCondition(access)
1173+
),
1174+
excludeSearchSources(excludedSources)
1175+
)}
11651176
), ranked_keyword_candidates AS MATERIALIZED (
1166-
SELECT * FROM scored_keyword_candidates
1167-
ORDER BY keyword_rank DESC, id LIMIT ${limit} OFFSET ${offset}
1177+
SELECT matched_keyword_chunks.id, matched_keyword_chunks.document_id,
1178+
${candidateRank} AS keyword_rank
1179+
FROM matched_keyword_chunks INNER JOIN ${embeddingKeywordSearch}
1180+
ON ${embeddingKeywordSearch.id} = matched_keyword_chunks.id
1181+
WHERE matched_keyword_chunks.document_id IN (SELECT id FROM visible_keyword_documents)
1182+
ORDER BY keyword_rank DESC, matched_keyword_chunks.id
1183+
LIMIT ${limit} OFFSET ${offset}
11681184
)
11691185
SELECT ranked_keyword_candidates.id, ${document.id} AS "documentId",
11701186
${document.connectorId} AS "connectorId",

0 commit comments

Comments
 (0)