From 18194c4bc4d648abea168d5b0b4355669f06d38d Mon Sep 17 00:00:00 2001 From: David Date: Mon, 8 Jun 2026 20:35:20 +0000 Subject: [PATCH] feat: emb_list and hybrid (RRF) searchIterator Adds Node SDK support for: - searchIterator() over an emb_list (ArrayOfVector) field. - hybridSearchIterator() fusing dense + emb_list modalities via RRF over the streamed per-modality results. Fusion keys on the primary key, which is force-included in each sub-iterator's output_fields so fusion is correct even when the caller omits the PK. Part of milvus-io/milvus#49906. Tests: test/utils/HybridSearchIterator.spec.ts. Signed-off-by: David --- examples/milvus/HybridSearchIterator.ts | 91 ++++++ milvus/grpc/Data.ts | 184 ++++++++++++- milvus/types/Search.ts | 12 + milvus/utils/HybridSearchIterator.ts | 252 +++++++++++++++++ milvus/utils/index.ts | 1 + test/utils/HybridSearchIterator.spec.ts | 351 ++++++++++++++++++++++++ 6 files changed, 889 insertions(+), 2 deletions(-) create mode 100644 examples/milvus/HybridSearchIterator.ts create mode 100644 milvus/utils/HybridSearchIterator.ts create mode 100644 test/utils/HybridSearchIterator.spec.ts diff --git a/examples/milvus/HybridSearchIterator.ts b/examples/milvus/HybridSearchIterator.ts new file mode 100644 index 00000000..359d48d2 --- /dev/null +++ b/examples/milvus/HybridSearchIterator.ts @@ -0,0 +1,91 @@ +import { MilvusClient, DataType } from '@zilliz/milvus2-sdk-node'; + +/** + * hybridSearchIterator: stream a hybrid search as RRF-fused batches. + * + * Each request in `data` is iterated as its own stateless single-modality + * search_iterator; the score-descending streams are fused client-side with + * Reciprocal Rank Fusion. `searchIterator` likewise accepts an emb_list + * (array-of-vector) query -- pass that query's multiple vectors as `data` + * against an emb_list `anns_field` (see the commented call below). + */ + +const COLLECTION_NAME = 'hybrid_search_iterator_demo'; + +(async () => { + const milvusClient = new MilvusClient({ + address: 'localhost:19530', + logLevel: 'info', + }); + + const dim = 8; + + await milvusClient.createCollection({ + collection_name: COLLECTION_NAME, + fields: [ + { + name: 'id', + data_type: DataType.Int64, + is_primary_key: true, + autoID: true, + }, + { name: 'title', data_type: DataType.VarChar, max_length: 128 }, + { name: 'dense_a', data_type: DataType.FloatVector, dim }, + { name: 'dense_b', data_type: DataType.FloatVector, dim }, + ], + }); + + const rows = 5000; + const data = []; + for (let i = 0; i < rows; i++) { + data.push({ + title: `doc_${i}`, + dense_a: new Array(dim).fill(0).map(() => Math.random()), + dense_b: new Array(dim).fill(0).map(() => Math.random()), + }); + } + await milvusClient.insert({ collection_name: COLLECTION_NAME, data }); + await milvusClient.flush({ collection_names: [COLLECTION_NAME] }); + + for (const field of ['dense_a', 'dense_b']) { + await milvusClient.createIndex({ + collection_name: COLLECTION_NAME, + field_name: field, + metric_type: 'COSINE', + }); + } + await milvusClient.loadCollectionSync({ collection_name: COLLECTION_NAME }); + + // --- hybrid search_iterator: RRF fusion of two modalities, streamed --------- + const hybrid = await milvusClient.hybridSearchIterator({ + collection_name: COLLECTION_NAME, + data: [ + { anns_field: 'dense_a', data: new Array(dim).fill(0).map(Math.random) }, + { anns_field: 'dense_b', data: new Array(dim).fill(0).map(Math.random) }, + ], + batchSize: 200, + rrf_k: 60, + output_fields: ['title'], + }); + + let page = 0; + for await (const batch of hybrid) { + page += 1; + // each batch is RRF-descending; `score` is the NRA lower bound on the RRF score + console.log(`hybrid page ${page}: ${batch.length} fused hits`, batch[0]); + } + + // --- emb_list search_iterator (shape reference) ---------------------------- + // For an emb_list (array-of-vector) field, pass one query's multiple vectors + // as `data` against the emb_list `anns_field`: + // + // const it = await milvusClient.searchIterator({ + // collection_name: COLLECTION_NAME, + // anns_field: 'paragraph_vectors', // an emb_list field + // data: [[/* vec1 */], [/* vec2 */], [/* vec3 */]], + // batchSize: 200, + // }); + // for await (const batch of it) { console.log(batch); } + + await milvusClient.dropCollection({ collection_name: COLLECTION_NAME }); +})(); diff --git a/milvus/grpc/Data.ts b/milvus/grpc/Data.ts index 9f18fbfc..28f7120d 100644 --- a/milvus/grpc/Data.ts +++ b/milvus/grpc/Data.ts @@ -43,7 +43,14 @@ import { SearchRes, SearchSimpleReq, SearchIteratorReq, + SearchResultData, HybridSearchReq, + HybridSearchIteratorReq, + RrfHybridFuser, + DEFAULT_RRF_K, + FetchBatch, + FusionItem, + FusionPk, promisify, sleep, parseToKeyValue, @@ -79,6 +86,7 @@ import { FloatVector, FieldPartialUpdateOpType, FieldPartialUpdateOp, + unixtimeToHybridts, } from '../'; import { Collection } from './Collection'; @@ -821,7 +829,9 @@ export class Data extends Collection { [ITERATOR_FIELD]: true, [ITER_SEARCH_V2_KEY]: true, [ITER_SEARCH_BATCH_SIZE_KEY]: batchSize, - [GUARANTEE_TIMESTAMP_KEY]: 0, + // use the caller's pinned snapshot if injected (hybridSearchIterator does + // this to share one snapshot across modalities); else self-pin below + [GUARANTEE_TIMESTAMP_KEY]: param.guarantee_timestamp ?? 0, [COLLECTION_ID]: collectionInfo.collectionID, }; @@ -849,7 +859,11 @@ export class Data extends Collection { batchRes.search_iterator_v2_results!.token; params[ITER_SEARCH_LAST_BOUND_KEY] = batchRes.search_iterator_v2_results?.last_bound; - params[GUARANTEE_TIMESTAMP_KEY] = batchRes.session_ts; + // self-pin the snapshot from the first batch, unless the caller + // injected a guarantee_timestamp (then keep that fixed) + if (!param.guarantee_timestamp) { + params[GUARANTEE_TIMESTAMP_KEY] = batchRes.session_ts; + } params[ITER_SEARCH_BATCH_SIZE_KEY] = batchSize; return { @@ -868,6 +882,172 @@ export class Data extends Collection { }; } + /** + * Executes a hybrid search and returns an async iterator over RRF-fused batches. + * + * Each request in `param.data` is iterated as its own stateless single-modality + * search_iterator; the score-descending streams are fused client-side with + * Reciprocal Rank Fusion via the NRA threshold algorithm (SPEC 6.6). One + * snapshot timestamp is pinned for the whole hybrid iterator and shared by + * every modality, so the fusion is over a single consistent snapshot. + * + * The emitted score is the NRA lower bound on the true RRF score (re-score if + * an exact RRF score is required). `output_fields` are forwarded to every + * modality and carried back through the fusion onto each result row. + * + * @param {HybridSearchIteratorReq} param - The hybrid search iterator request. + * @returns {Promise} - An async iterator yielding RRF-fused result batches. + * + * @example + * const iterator = await client.hybridSearchIterator({ + * collection_name: 'my_collection', + * data: [ + * { anns_field: 'vector', data: [0.1, 0.2, 0.3, 0.4] }, + * { anns_field: 'emb_list', data: [[0.1, 0.2], [0.3, 0.4]] }, // emb_list query + * ], + * batchSize: 100, + * output_fields: ['title'], + * }); + * + * for await (const batch of iterator) { + * console.log(batch); // each batch is RRF-descending + * } + */ + async hybridSearchIterator(param: HybridSearchIteratorReq): Promise { + // batchSize lower-bound validation (the high end is capped per modality by + // searchIterator); parity with pymilvus, which rejects a non-positive batch_size + if (typeof param.batchSize !== 'number' || param.batchSize < 1) { + throw new Error('batchSize must be a positive integer'); + } + + const client = this; + const rrfK = param.rrf_k ?? DEFAULT_RRF_K; + const limit = + !param.limit || param.limit === NO_LIMIT ? Infinity : param.limit; + + // Pin ONE snapshot for the whole hybrid iterator and share it with every + // modality, so the fused result is over a single consistent snapshot -- + // mirroring pymilvus's HybridSearchIteratorV2. Each modality below is a real + // single-modality searchIterator (stateless Iterator v2) composed here; the + // injected guarantee_timestamp is what makes the shared snapshot possible -- + // without it each searchIterator self-pins its own per-instance snapshot. + const pinnedTs = + param.guarantee_timestamp ?? + unixtimeToHybridts(Math.floor(Date.now() / 1000).toString()); + + // each emitted doc's hit, kept so next() can carry output_fields through the + // fusion; dropped on emit and pruned against the fuser's emitted set + const hitsById = new Map(); + + // The RRF fuser keys each modality's hits on the primary key (hit.id). The + // sdk's formatSearchResult only populates hit.id when the PK field is present + // in output_fields -- otherwise every hit has id === undefined and they all + // collapse to a single fusion key. So force-include the PK field in every + // sub-iterator's output_fields (it's the join key for fusion). + const pkName = await client.getPkFieldName({ + collection_name: param.collection_name, + db_name: param.db_name, + } as DescribeCollectionReq); + const subOutputFields = + param.output_fields && param.output_fields.includes(pkName) + ? param.output_fields + : [...(param.output_fields ?? []), pkName]; + + // compose one real single-modality searchIterator per request, all sharing + // the pinned snapshot ts + const subIterators = await Promise.all( + param.data.map(req => + client.searchIterator({ + collection_name: param.collection_name, + db_name: param.db_name, + anns_field: req.anns_field, + data: req.data, + expr: req.expr, + exprValues: req.exprValues, + output_fields: subOutputFields, + consistency_level: param.consistency_level, + ignore_growing: req.ignore_growing, + group_by_field: req.group_by_field, + params: req.params, + batchSize: param.batchSize, + guarantee_timestamp: pinnedTs, + }) + ) + ); + + // adapt each searchIterator into a () -> [(id, score)] batch source for the + // RRF fuser, stashing each hit so next() can return its output_fields + const fetchBatches: FetchBatch[] = subIterators.map(subIt => { + const cursor = subIt[Symbol.asyncIterator](); + return async (): Promise => { + const { value } = await cursor.next(); + const hits = (value as SearchResultData[]) || []; + const items: FusionItem[] = []; + for (const hit of hits) { + // Defensive: a missing PK would silently collapse all hits onto one + // fusion key (see the output_fields note above). Fail loud instead. + if (hit.id === undefined || hit.id === null) { + throw new Error( + 'hybridSearchIterator: a hit is missing its primary key (id); ' + + 'cannot fuse modalities. The collection must have a primary key.' + ); + } + hitsById.set(hit.id, hit); + items.push([hit.id, Number(hit.score)]); + } + return items; + }; + }); + + const fuser = new RrfHybridFuser(fetchBatches, rrfK); + let currentTotal = 0; + + return { + [Symbol.asyncIterator]() { + return { + async next() { + if (currentTotal >= limit) { + return { done: true, value: null }; + } + const target = Math.min(param.batchSize, limit - currentTotal); + + let fused: FusionItem[]; + try { + fused = await fuser.nextBatch(target); + } catch (error) { + console.error('Error during hybrid search iteration:', error); + return { done: true, value: null }; + } + + if (fused.length === 0) { + return { done: true, value: [] }; + } + currentTotal += fused.length; + + // carry output_fields back through the fusion: take each emitted + // doc's stashed hit and replace its score with the fused RRF score + const value: SearchResultData[] = fused.map(([id, score]) => { + const hit = hitsById.get(id); + hitsById.delete(id); + return { ...(hit as SearchResultData), score }; + }); + + // a doc can resurface in another modality after it was emitted: the + // fuser skips it, but its hit was still stashed -- drop those so + // hitsById stays bounded by dense/sparse stream skew + for (const key of [...hitsById.keys()]) { + if (fuser.hasEmitted(key)) { + hitsById.delete(key); + } + } + + return { done: false, value }; + }, + }; + }, + }; + } + /** * Executes a query and returns an async iterator that allows iterating over the results in batches. * diff --git a/milvus/types/Search.ts b/milvus/types/Search.ts index eb791b57..211ee013 100644 --- a/milvus/types/Search.ts +++ b/milvus/types/Search.ts @@ -141,6 +141,7 @@ export interface SearchIteratorReq limit?: number; // Optional. Specifies the maximum number of items. Default is no limit (-1 or if not set). batchSize: number; // Specifies the number of items to return in each batch. if it exceeds 16384, it will be set to 16384 external_filter_fn?: (row: SearchResultData) => boolean; // Optional. Specifies the external filter function. + guarantee_timestamp?: string | number; // Optional. A pinned snapshot timestamp; when set, the iterator reads this fixed snapshot instead of self-pinning from the first batch. hybridSearchIterator injects it to share one snapshot across modalities. } // rerank strategy and parameters @@ -160,6 +161,17 @@ export type HybridSearchReq = Omit< rerank?: RerankerObj | FunctionObject | FunctionScore; }; +// hybrid search_iterator parameter type. Each request in `data` is iterated as +// its own stateless single-modality search_iterator; the score-descending +// streams are fused client-side with Reciprocal Rank Fusion (SPEC 6.6). The +// server-side `rerank` is replaced by client-side RRF, tuned by `rrf_k`. +export type HybridSearchIteratorReq = Omit & { + batchSize: number; // fused results per batch; capped at DEFAULT_MAX_SEARCH_SIZE + limit?: number; // total fused results to return; default no limit + rrf_k?: number; // RRF constant k; defaults to 60 + guarantee_timestamp?: string | number; // Optional. A pinned snapshot timestamp shared by every modality; defaults to one pinned at call time. +}; + // search api response type export interface SearchRes extends resStatusResponse { results: { diff --git a/milvus/utils/HybridSearchIterator.ts b/milvus/utils/HybridSearchIterator.ts new file mode 100644 index 00000000..9d92c07e --- /dev/null +++ b/milvus/utils/HybridSearchIterator.ts @@ -0,0 +1,252 @@ +/** + * Client-side RRF fusion for the hybrid search_iterator. + * + * Per SPEC 6.6 (emb_list iterator series), hybrid iteration is done in the SDK: + * each modality is a stateless single-field search_iterator, and the + * score-descending streams are fused with Reciprocal Rank Fusion via the NRA + * threshold algorithm. This is the TypeScript port of pymilvus's + * `_RrfHybridFuser` / `_StreamCursor` (PR 6a) -- the async batch sources are the + * only adaptation; the fusion logic is identical. + * + * References: Fagin/Lotem/Naor 2003 (NRA threshold algorithm); + * Cormack et al. 2009 (Reciprocal Rank Fusion). + */ + +/** Reciprocal-rank-fusion constant `k` (Cormack et al. 2009). */ +export const DEFAULT_RRF_K = 60; + +/** + * Soft bound on the NRA in-flight (seen-but-not-emitted) map. The map is bounded + * by dense/sparse stream skew; on overflow the fuser force-flushes its best + * in-flight document rather than failing the iteration -- client-side memory is + * cheap. + */ +export const DEFAULT_INFLIGHT_CAP = 100000; + +/** A document primary key. */ +export type FusionPk = string | number; + +/** One stream item: a primary key and its (descending) score. */ +export type FusionItem = [FusionPk, number]; + +/** + * A batch source: resolves to the next score-descending batch of items, or an + * empty array once the stream is exhausted. + */ +export type FetchBatch = () => Promise; + +/** + * Item-by-item, lazily-refilled view over one batched, score-descending stream. + */ +export class StreamCursor { + private buf: FusionItem[] = []; + private pos = 0; + /** Number of items consumed so far. */ + reads = 0; + /** True once the underlying stream has yielded its final (empty) batch. */ + exhausted = false; + + constructor(private readonly fetchBatch: FetchBatch) {} + + /** Consume and return the next item, or null when the stream is exhausted. */ + async advance(): Promise { + if (this.pos >= this.buf.length) { + if (this.exhausted) { + return null; + } + this.buf = (await this.fetchBatch()) || []; + this.pos = 0; + if (this.buf.length === 0) { + this.exhausted = true; + return null; + } + } + const item = this.buf[this.pos]; + this.pos += 1; + this.reads += 1; + return item; + } +} + +/** + * Incremental RRF fusion of N score-descending streams via the NRA algorithm. + * + * RRF is rank-based: a stream's contribution to any not-yet-seen document is + * `1 / (k + reads_so_far + 1)`, independent of raw scores. The fuser advances + * the least-read live stream and settles a document once no in-flight or unseen + * document can overtake it. + * + * Emission order is exact: each emitted document provably out-ranks every + * not-yet-emitted document by true RRF score. The emitted *score*, however, is a + * lower bound -- a document is settled with the partial RRF score of the streams + * it has been seen in so far; if it later resurfaces in a not-yet-consumed + * stream, it is already emitted and the extra rank is dropped. Callers that need + * the exact RRF score must re-score. + */ +export class RrfHybridFuser { + private readonly cursors: StreamCursor[]; + /** doc id -> (stream index -> rank); the NRA in-flight (seen-but-not-emitted) map. */ + private readonly seen = new Map>(); + /** + * Doc ids already emitted -- a stream resurfacing one must not re-add it to + * `seen` (that would re-emit it; deep skewed streams hit this readily). + */ + private readonly emitted = new Set(); + + constructor( + fetchBatches: FetchBatch[], + private readonly rrfK: number = DEFAULT_RRF_K, + private readonly inflightCap: number = DEFAULT_INFLIGHT_CAP + ) { + this.cursors = fetchBatches.map(fb => new StreamCursor(fb)); + } + + /** Whether `id` has already been emitted (used to prune callers' hit maps). */ + hasEmitted(id: FusionPk): boolean { + return this.emitted.has(id); + } + + /** An exhausted stream can never produce another rank, so it contributes 0. */ + private contribution(cursor: StreamCursor): number { + return cursor.exhausted ? 0 : 1 / (this.rrfK + cursor.reads + 1); + } + + /** RRF score from the ranks already known for a document. */ + private worst(ranks: Map): number { + let sum = 0; + for (const rank of ranks.values()) { + sum += 1 / (this.rrfK + rank); + } + return sum; + } + + /** Advance the least-read live stream by one item (largest RRF contribution). */ + private async advanceLeastRead(): Promise { + let idx = -1; + let cursor: StreamCursor | null = null; + this.cursors.forEach((c, i) => { + if (c.exhausted) { + return; + } + if (cursor === null || c.reads < cursor.reads) { + idx = i; + cursor = c; + } + }); + if (cursor === null) { + return; + } + const item = await (cursor as StreamCursor).advance(); + if (item !== null) { + const docId = item[0]; + // an already-emitted doc resurfacing in another stream is dropped -- + // it is ranked; re-adding it to `seen` would emit a duplicate + if (!this.emitted.has(docId)) { + let ranks = this.seen.get(docId); + if (!ranks) { + ranks = new Map(); + this.seen.set(docId, ranks); + } + if (!ranks.has(idx)) { + ranks.set(idx, (cursor as StreamCursor).reads); + } + } + } + } + + /** + * Return the best in-flight document if it is provably the next global + * result, else null. + */ + private settled(): FusionItem | null { + if (this.seen.size === 0) { + return null; + } + let tau = 0; + for (const c of this.cursors) { + tau += this.contribution(c); + } + const worsts = new Map(); + for (const [doc, ranks] of this.seen) { + worsts.set(doc, this.worst(ranks)); + } + // candidate = the in-flight doc with the highest known (worst-case) score + let cand: FusionPk | null = null; + let candWorst = -Infinity; + for (const [doc, w] of worsts) { + if (w > candWorst) { + candWorst = w; + cand = doc; + } + } + if (cand === null) { + return null; + } + // an unseen document could still outrank cand + if (candWorst < tau) { + return null; + } + // an in-flight document could still outrank cand + for (const [doc, ranks] of this.seen) { + if (doc === cand) { + continue; + } + let best = worsts.get(doc) as number; + for (let i = 0; i < this.cursors.length; i++) { + if (!ranks.has(i)) { + best += this.contribution(this.cursors[i]); + } + } + if (best > candWorst) { + return null; + } + } + return [cand, candWorst]; + } + + private emit(out: FusionItem[], docId: FusionPk, score: number): void { + this.seen.delete(docId); + this.emitted.add(docId); + out.push([docId, score]); + } + + /** + * Fuse and return up to `batchSize` (docId, rrfScore) items. + * + * Emitted RRF-descending; the score is a lower bound on the true RRF score + * (see the class docstring). State persists across calls for the lifetime of + * the iterator. + */ + async nextBatch(batchSize: number): Promise { + const out: FusionItem[] = []; + while (out.length < batchSize) { + const settled = this.settled(); + if (settled !== null) { + this.emit(out, settled[0], settled[1]); + continue; + } + if (this.cursors.every(c => c.exhausted)) { + // no stream can advance; settled() has already drained `seen` + break; + } + await this.advanceLeastRead(); + // force-flush the best in-flight doc on overflow -- bound memory on + // skewed streams without failing the iteration + if (this.inflightCap && this.seen.size > this.inflightCap) { + let cand: FusionPk | null = null; + let candWorst = -Infinity; + for (const [doc, ranks] of this.seen) { + const w = this.worst(ranks); + if (w > candWorst) { + candWorst = w; + cand = doc; + } + } + if (cand !== null) { + this.emit(out, cand, candWorst); + } + } + } + return out; + } +} diff --git a/milvus/utils/index.ts b/milvus/utils/index.ts index 161e4014..5ba65c81 100644 --- a/milvus/utils/index.ts +++ b/milvus/utils/index.ts @@ -3,6 +3,7 @@ export * from './Connection'; export * from './Schema'; export * from './Data'; export * from './Search'; +export * from './HybridSearchIterator'; export * from './Bytes'; export * from './Format'; export * from './Validate'; diff --git a/test/utils/HybridSearchIterator.spec.ts b/test/utils/HybridSearchIterator.spec.ts new file mode 100644 index 00000000..87cded70 --- /dev/null +++ b/test/utils/HybridSearchIterator.spec.ts @@ -0,0 +1,351 @@ +import { + RrfHybridFuser, + StreamCursor, + FetchBatch, + FusionItem, + FusionPk, +} from '../../milvus'; +import { Data } from '../../milvus/grpc/Data'; + +/** + * Unit tests for the client-side hybrid searchIterator RRF fusion (PR 6b) -- + * the TypeScript port of pymilvus's _RrfHybridFuser. These exercise the NRA/RRF + * fuser and the hybridSearchIterator output_fields passthrough over mock + * streams / a fake client -- no live milvus. + */ + +const K = 60; + +// a () -> Promise batch source serving `items` in fixed batches +function makeSource(items: FusionItem[], batchSize = 4): FetchBatch { + const batches: FusionItem[][] = []; + for (let i = 0; i < items.length; i += batchSize) { + batches.push(items.slice(i, i + batchSize)); + } + let idx = 0; + return async () => (idx < batches.length ? batches[idx++] : []); +} + +// brute-force RRF: score(d) = sum over streams where d appears of 1/(k+rank) +function referenceRrf( + streams: FusionItem[][], + k: number +): Map { + const ranks = new Map>(); + streams.forEach((stream, si) => { + stream.forEach(([doc], i) => { + if (!ranks.has(doc)) ranks.set(doc, new Map()); + ranks.get(doc)!.set(si, i + 1); + }); + }); + const score = new Map(); + for (const [doc, byStream] of ranks) { + let s = 0; + for (const r of byStream.values()) s += 1 / (k + r); + score.set(doc, s); + } + return score; +} + +async function drain( + fuser: RrfHybridFuser, + batchSize: number +): Promise { + const out: FusionItem[] = []; + for (;;) { + const batch = await fuser.nextBatch(batchSize); + if (batch.length === 0) break; + out.push(...batch); + } + return out; +} + +// `fused` is a correct RRF fusion of `streams`: complete, deduplicated, emitted +// in true-RRF-descending order, each emitted score a (non-increasing) lower +// bound on the true RRF score. See RrfHybridFuser -- the NRA emission rule. +function assertValidRrf( + fused: FusionItem[], + streams: FusionItem[][], + k: number +) { + const ref = referenceRrf(streams, k); + const ids = fused.map(([d]) => d); + expect(ids.length).toBe(new Set(ids).size); // no duplicate + expect(new Set(ids)).toEqual(new Set(ref.keys())); // complete + + const trueSeq = fused.map(([d]) => ref.get(d)!); + for (let i = 1; i < trueSeq.length; i++) { + expect(trueSeq[i - 1]).toBeGreaterThanOrEqual(trueSeq[i] - 1e-9); + } + const emitted = fused.map(([, s]) => s); + for (let i = 1; i < emitted.length; i++) { + expect(emitted[i - 1]).toBeGreaterThanOrEqual(emitted[i] - 1e-9); + } + for (const [d, s] of fused) { + expect(s).toBeLessThanOrEqual(ref.get(d)! + 1e-9); + } +} + +describe('RrfHybridFuser', () => { + it('fuses disjoint streams', async () => { + const s0: FusionItem[] = Array.from({ length: 10 }, (_, i) => [ + `a${i}`, + 1 - i * 0.01, + ]); + const s1: FusionItem[] = Array.from({ length: 10 }, (_, i) => [ + `b${i}`, + 1 - i * 0.01, + ]); + const fuser = new RrfHybridFuser([makeSource(s0), makeSource(s1)], K); + assertValidRrf(await drain(fuser, 5), [s0, s1], K); + }); + + it('sums both ranks for an overlapping doc', async () => { + // "x" is rank 1 in s0 and rank 1 in s1 -> highest RRF (2/(K+1)) + const s0: FusionItem[] = [ + ['x', 0.9], + ['a', 0.8], + ['b', 0.7], + ]; + const s1: FusionItem[] = [ + ['x', 0.5], + ['c', 0.4], + ['d', 0.3], + ]; + const fuser = new RrfHybridFuser([makeSource(s0), makeSource(s1)], K); + const fused = await drain(fuser, 10); + assertValidRrf(fused, [s0, s1], K); + expect(fused[0][0]).toBe('x'); + expect(fused[0][1]).toBeCloseTo(2 / (K + 1), 9); + }); + + it('handles full overlap with reversed orders', async () => { + const s0: FusionItem[] = Array.from({ length: 8 }, (_, i) => [ + `d${i}`, + 1 - i * 0.1, + ]); + const s1: FusionItem[] = [...s0].reverse(); + const fuser = new RrfHybridFuser([makeSource(s0), makeSource(s1)], K); + const fused = await drain(fuser, 3); + assertValidRrf(fused, [s0, s1], K); + expect(fused.length).toBe(8); + }); + + it('emits a deep skewed shared doc exactly once', async () => { + // `shared` settles from its rank-1 sighting in the single-item stream s0 + // long before the deep stream s1 reaches it at rank 51 -- the resurfacing + // must not re-emit it (the emitted-set guard). + const shared = 'shared'; + const s0: FusionItem[] = [[shared, 1.0]]; + const s1: FusionItem[] = [ + ...Array.from( + { length: 50 }, + (_, i): FusionItem => [`b${i}`, 0.9 - i * 0.001] + ), + [shared, 0.5], + ...Array.from( + { length: 40 }, + (_, i): FusionItem => [`b${i + 50}`, 0.4 - i * 0.001] + ), + ]; + const fuser = new RrfHybridFuser([makeSource(s0), makeSource(s1)], K); + const fused = await drain(fuser, 5); + const ids = fused.map(([d]) => d); + expect(ids.filter(d => d === shared).length).toBe(1); + expect(ids.length).toBe(91); // shared + 90 distinct b's + // settled from the rank-1 sighting only -> emitted score is the lower bound + const sharedScore = fused.find(([d]) => d === shared)![1]; + expect(sharedScore).toBeCloseTo(1 / (K + 1), 9); + expect(sharedScore).toBeLessThan(1 / (K + 1) + 1 / (K + 51)); + assertValidRrf(fused, [s0, s1], K); + }); + + it('handles one empty stream', async () => { + const s0: FusionItem[] = Array.from({ length: 12 }, (_, i) => [ + `a${i}`, + 1 - i * 0.01, + ]); + const fuser = new RrfHybridFuser([makeSource(s0), makeSource([])], K); + assertValidRrf(await drain(fuser, 5), [s0, []], K); + }); + + it('handles two empty streams', async () => { + const fuser = new RrfHybridFuser([makeSource([]), makeSource([])], K); + expect(await fuser.nextBatch(10)).toEqual([]); + }); + + it('is consistent across batch sizes', async () => { + const s0: FusionItem[] = Array.from({ length: 20 }, (_, i) => [ + `a${i}`, + 1 - i * 0.01, + ]); + const s1: FusionItem[] = Array.from({ length: 20 }, (_, i) => [ + `b${i}`, + 1 - i * 0.01, + ]); + s1[5] = ['a3', 0.5]; + s1[9] = ['a7', 0.4]; + const oneShot = await drain( + new RrfHybridFuser([makeSource(s0), makeSource(s1)], K), + 10000 + ); + const small = await drain( + new RrfHybridFuser([makeSource(s0), makeSource(s1)], K), + 3 + ); + expect(oneShot).toEqual(small); + }); + + it('force-flushes under a tiny in-flight cap, still yielding every doc once', async () => { + const s0: FusionItem[] = Array.from({ length: 40 }, (_, i) => [ + `a${i}`, + 1 - i * 0.001, + ]); + const s1: FusionItem[] = Array.from({ length: 40 }, (_, i) => [ + `b${i}`, + 1 - i * 0.001, + ]); + const fuser = new RrfHybridFuser([makeSource(s0), makeSource(s1)], K, 4); + const ids = (await drain(fuser, 7)).map(([d]) => d); + expect(ids.length).toBe(80); + expect(new Set(ids).size).toBe(80); + }); + + it('reports emitted docs via hasEmitted', async () => { + const s0: FusionItem[] = [ + ['a', 0.9], + ['b', 0.8], + ]; + const fuser = new RrfHybridFuser([makeSource(s0), makeSource([])], K); + expect(fuser.hasEmitted('a')).toBe(false); + await drain(fuser, 10); + expect(fuser.hasEmitted('a')).toBe(true); + expect(fuser.hasEmitted('b')).toBe(true); + expect(fuser.hasEmitted('missing')).toBe(false); + }); +}); + +describe('StreamCursor', () => { + it('lazily refills and reports exhaustion', async () => { + const items: FusionItem[] = Array.from({ length: 5 }, (_, i) => [i, 1.0]); + const cursor = new StreamCursor(makeSource(items, 2)); + const got: FusionPk[] = []; + for (;;) { + const item = await cursor.advance(); + if (item === null) break; + got.push(item[0]); + } + expect(got).toEqual([0, 1, 2, 3, 4]); + expect(cursor.reads).toBe(5); + expect(cursor.exhausted).toBe(true); + expect(await cursor.advance()).toBeNull(); // stays exhausted + }); +}); + +// --- hybridSearchIterator wrapper: output_fields passthrough -------------------- + +// a fake Data client: hybridSearchIterator composes real searchIterator +// instances, so the fake stubs searchIterator -- each returns an async-iterable +// replaying the pages handed in for that anns_field. No live milvus. +function fakeClient(streamsByField: Record) { + return { + // the fix force-includes the PK field in each sub-iterator's output_fields + // so hit.id is populated as the fusion key; the fake hits key on `id`. + getPkFieldName: async () => 'id', + searchIterator: async (req: any) => { + const pages = streamsByField[req.anns_field] || []; + let idx = 0; + return { + [Symbol.asyncIterator]() { + return { + async next() { + if (idx < pages.length) { + return { done: false, value: pages[idx++] }; + } + return { done: true, value: null }; + }, + }; + }, + }; + }, + }; +} + +async function drainIterator(it: any): Promise { + const out: any[] = []; + for await (const batch of it) { + if (batch && batch.length) out.push(...batch); + } + return out; +} + +describe('hybridSearchIterator', () => { + it('carries output_fields through the fusion and replaces score with RRF', async () => { + const dense = [ + [ + { id: 'x', score: 0.9, title: 'doc x' }, + { id: 'a', score: 0.8, title: 'doc a' }, + ], + ]; + const sparse = [ + [ + { id: 'x', score: 0.5, title: 'doc x' }, + { id: 'c', score: 0.4, title: 'doc c' }, + ], + ]; + const fake = fakeClient({ dense, sparse }); + const it = await (Data.prototype.hybridSearchIterator as any).call(fake, { + collection_name: 'c', + data: [ + { anns_field: 'dense', data: [0.1, 0.2] }, + { anns_field: 'sparse', data: [0.3, 0.4] }, + ], + batchSize: 10, + output_fields: ['title'], + }); + const rows = await drainIterator(it); + const ids = rows.map(r => r.id); + expect(new Set(ids)).toEqual(new Set(['x', 'a', 'c'])); + // every fused row keeps its entity field + for (const row of rows) { + expect(typeof row.title).toBe('string'); + expect(row.title).toContain('doc'); + } + // "x" is rank 1 in both modalities -> top result, RRF score 2/(K+1) + expect(rows[0].id).toBe('x'); + expect(rows[0].score).toBeCloseTo(2 / (K + 1), 9); + // a doc shared by both modalities is emitted once + expect(ids.filter(d => d === 'x').length).toBe(1); + }); + + it('respects the limit', async () => { + const dense = [ + Array.from({ length: 20 }, (_, i) => ({ + id: `a${i}`, + score: 1 - i * 0.01, + })), + ]; + const fake = fakeClient({ dense }); + const it = await (Data.prototype.hybridSearchIterator as any).call(fake, { + collection_name: 'c', + data: [{ anns_field: 'dense', data: [0.1] }], + batchSize: 5, + limit: 7, + }); + const rows = await drainIterator(it); + expect(rows.length).toBe(7); + }); + + it('rejects a non-positive batchSize', async () => { + const fake = fakeClient({ dense: [] }); + for (const batchSize of [0, -1]) { + await expect( + (Data.prototype.hybridSearchIterator as any).call(fake, { + collection_name: 'c', + data: [{ anns_field: 'dense', data: [0.1] }], + batchSize, + }) + ).rejects.toThrow('batchSize must be a positive integer'); + } + }); +});