From 6fa92e3bf15b1f9238f2c7faf0810ebb1361953a Mon Sep 17 00:00:00 2001 From: shanghaikid Date: Fri, 10 Jan 2025 14:18:26 +0800 Subject: [PATCH 1/4] support search iterator Signed-off-by: shanghaikid --- milvus/grpc/Data.ts | 270 +++++++++---------------- milvus/types/Data.ts | 34 ++-- milvus/utils/Format.ts | 8 +- milvus/utils/Function.ts | 21 -- test/grpc/Iterator.spec.ts | 386 +++++++++++++++++------------------- test/utils/Format.spec.ts | 11 +- test/utils/Function.spec.ts | 133 ------------- 7 files changed, 305 insertions(+), 558 deletions(-) diff --git a/milvus/grpc/Data.ts b/milvus/grpc/Data.ts index db77bae3..6bf1838b 100644 --- a/milvus/grpc/Data.ts +++ b/milvus/grpc/Data.ts @@ -37,7 +37,6 @@ import { SearchRes, SearchSimpleReq, SearchIteratorReq, - DEFAULT_TOPK, HybridSearchReq, promisify, sleep, @@ -56,9 +55,6 @@ import { DEFAULT_COUNT_QUERY_STRING, getQueryIteratorExpr, QueryIteratorReq, - getRangeFromSearchResult, - SearchResultData, - getPKFieldExpr, DEFAULT_MAX_SEARCH_SIZE, SparseFloatVector, sparseRowsToBytes, @@ -473,6 +469,10 @@ export class Data extends Collection { * @returns {string} status.error_code - The error code of the operation. * @returns {string} status.reason - The reason for the error, if any. * @returns {{score:number,id:string, [outputfield]: value}[]} results - Array of search results. + * @returns {number} session_ts - The timestamp of the search session. + * @returns {string} collection_name - The name of the collection. + * @returns {number} all_search_count - The total number of search operations. + * @returns {string[]} recalls - The recalls of the search operation. * * @example * ``` @@ -533,6 +533,8 @@ export class Data extends Collection { status: originSearchResult.status, results: [], recalls: [], + session_ts: -1, + collection_name: data.collection_name, }; } @@ -547,183 +549,95 @@ export class Data extends Collection { // nq === 1, return the first object of results array results: nq === 1 ? results[0] || [] : results, recalls: originSearchResult.results.recalls, + session_ts: originSearchResult.session_ts, + collection_name: data.collection_name, + all_search_count: originSearchResult.results.all_search_count, + search_iterator_v2_results: + originSearchResult.results.search_iterator_v2_results, + _search_iterator_v2_results: + originSearchResult.results._search_iterator_v2_results, }; } - // async searchIterator(data: SearchIteratorReq): Promise { - // // store client - // const client = this; - // // get collection info - // const pkField = await this.getPkField(data); - // // get available count - // const count = await client.count({ - // collection_name: data.collection_name, - // expr: data.expr || data.filter || '', - // }); - // // make sure limit is not exceed the total count - // const total = data.limit > count.data ? count.data : data.limit; - // // make sure batch size is exceed the total count - // let batchSize = data.batchSize > total ? total : data.batchSize; - // // make sure batch size is not exceed max search size - // batchSize = - // batchSize > DEFAULT_MAX_SEARCH_SIZE ? DEFAULT_MAX_SEARCH_SIZE : batchSize; - - // // init expr - // const initExpr = data.expr || data.filter || ''; - // // init search params object - // data.params = data.params || {}; - // data.limit = batchSize; - - // // user range filter set - // const initRadius = Number(data.params.radius) || 0; - // const initRangeFilter = Number(data.params.range_filter) || 0; - // // range params object - // const rangeFilterParams = { - // radius: initRadius, - // rangeFilter: initRangeFilter, - // expr: initExpr, - // }; - - // // force quite if true, at first, if total is 0, return done - // let done = total === 0; - // // batch result store - // let lastBatchRes: SearchResultData[] = []; - - // // build cache - // const cache = await client.search({ - // ...data, - // limit: total > DEFAULT_MAX_SEARCH_SIZE ? DEFAULT_MAX_SEARCH_SIZE : total, - // }); - - // return { - // currentTotal: 0, - // [Symbol.asyncIterator]() { - // return { - // currentTotal: this.currentTotal, - // async next() { - // // check if reach the limit - // if ( - // (this.currentTotal >= total && this.currentTotal !== 0) || - // done - // ) { - // return { done: true, value: lastBatchRes }; - // } - - // // batch result container - // const batchRes: SearchResultData[] = []; - // const bs = - // this.currentTotal + batchSize > total - // ? total - this.currentTotal - // : batchSize; - - // // keep getting search data if not reach the batch size - // while (batchRes.length < bs) { - // // search results container - // let searchResults: SearchResults = { - // status: { error_code: 'SUCCESS', reason: '' }, - // results: [], - // }; - - // // Iterate through the cached data, adding it to the search results container until the batch size is reached. - // if (cache.results.length > 0) { - // while ( - // cache.results.length > 0 && - // searchResults.results.length < bs - // ) { - // searchResults.results.push(cache.results.shift()!); - // } - // } else if (searchResults.results.length < bs) { - // // build search params, overwrite range filter - // if (rangeFilterParams.radius && rangeFilterParams.rangeFilter) { - // data.params = { - // ...data.params, - // radius: rangeFilterParams.radius, - // range_filter: - // rangeFilterParams.rangeFilter, - // }; - // } - // // set search expr - // data.expr = rangeFilterParams.expr; - - // console.log('search param', data.params, data.expr); - - // // iterate search, if no result, double the radius, until we doubled for 5 times - // let newSearchRes = await client.search(data); - // let retry = 0; - // while (newSearchRes.results.length === 0 && retry < 5) { - // newSearchRes = await client.search(data); - // if (searchResults.results.length === 0) { - // const newRadius = rangeFilterParams.radius * 2; - - // data.params = { - // ...data.params, - // radius: newRadius, - // }; - // } - - // retry++; - // } - - // // combine search results - // searchResults.results = [ - // ...searchResults.results, - // ...newSearchRes.results, - // ]; - // } - - // console.log('return', searchResults.results); - - // // filter result, batchRes should be unique - // const filterResult = searchResults.results.filter( - // r => - // !lastBatchRes.some(l => l.id === r.id) && - // !batchRes.some(c => c.id === r.id) - // ); - - // // fill filter result to batch result, it should not exceed the batch size - // for (let i = 0; i < filterResult.length; i++) { - // if (batchRes.length < bs) { - // batchRes.push(filterResult[i]); - // } - // } - - // // get data range about last batch result - // const resultRange = getRangeFromSearchResult(filterResult); - - // console.log('result range', resultRange); - - // // if no more result, force quite - // if (resultRange.lastDistance === 0) { - // done = true; - // return { done: false, value: batchRes }; - // } - - // // update next range and expr - // rangeFilterParams.rangeFilter = resultRange.lastDistance; - // rangeFilterParams.radius = - // rangeFilterParams.radius + resultRange.radius; - // rangeFilterParams.expr = getPKFieldExpr({ - // pkField, - // value: resultRange.id as string, - // expr: initExpr, - // }); - - // console.log('last', rangeFilterParams); - // } - - // // store last result - // lastBatchRes = batchRes; - - // // update current total - // this.currentTotal += batchRes.length; - - // // return batch result - // return { done: false, value: batchRes }; - // }, - // }; - // }, - // }; - // } + async searchIterator(data: SearchIteratorReq): Promise { + const client = this; + + // Get available count + const count = await client.count({ + collection_name: data.collection_name, + expr: data.expr || data.filter || '', + }); + + // if limit not set, set it to count + if (!data.limit || data.limit === NO_LIMIT) { + data.limit = count.data; + } + + // Ensure limit does not exceed the total count + const total = Math.min(data.limit, count.data); + + // Ensure batch size does not exceed the total count or max search size + let batchSize = Math.min(data.batchSize, total, DEFAULT_MAX_SEARCH_SIZE); + + // Iterator fields + const ITERATOR_FIELD = 'iterator'; + const ITER_SEARCH_V2_KEY = 'search_iter_v2'; + const ITER_SEARCH_ID_KEY = 'search_iter_id'; + const ITER_SEARCH_BATCH_SIZE_KEY = 'search_iter_batch_size'; + const ITER_SEARCH_LAST_BOUND_KEY = 'search_iter_last_bound'; + const GUARANTEE_TIMESTAMP_KEY = 'guarantee_timestamp'; + + let currentTotal = 0; + + // search iterator special params + const params: any = { + ...data.params, + [ITERATOR_FIELD]: true, + [ITER_SEARCH_V2_KEY]: true, + [ITER_SEARCH_BATCH_SIZE_KEY]: batchSize, + [GUARANTEE_TIMESTAMP_KEY]: 0, + }; + + return { + [Symbol.asyncIterator]() { + return { + async next() { + if (currentTotal >= total) { + return { done: true, value: null }; + } + + try { + const batchRes = await client.search({ + ...data, + params, + limit: batchSize, + }); + + // update current total and batch size + currentTotal += batchRes.results.length; + batchSize = Math.min(batchSize, total - currentTotal); + + // update search params + params[ITER_SEARCH_ID_KEY] = + 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 || 0; + params[ITER_SEARCH_BATCH_SIZE_KEY] = batchSize; + + return { + done: currentTotal > total, + value: batchRes.results, + }; + } catch (error) { + console.error('Error during search iteration:', error); + return { done: true, value: null }; + } + }, + }; + }, + }; + } /** * Executes a query and returns an async iterator that allows iterating over the results in batches. diff --git a/milvus/types/Data.ts b/milvus/types/Data.ts index 8675a54e..9ed544ad 100644 --- a/milvus/types/Data.ts +++ b/milvus/types/Data.ts @@ -1,4 +1,3 @@ -import { Key } from 'readline'; import { GrpcTimeOut, KeyValuePair, @@ -263,6 +262,11 @@ export interface SearchResultData { export interface SearchResults extends resStatusResponse { results: SearchResultData[]; recalls: number[]; + session_ts: number; + collection_name: string; + all_search_count?: number; + search_iterator_v2_results?: Record; + _search_iterator_v2_results?: string; } export interface ImportResponse extends resStatusResponse { @@ -302,6 +306,7 @@ export interface SearchParam { group_size?: number; // group size strict_group_size?: boolean; // if strict group size hints?: string; // hints to improve milvus search performance + [key: string]: any; // extra search parameters } // old search api parameter type, deprecated @@ -320,16 +325,6 @@ export interface SearchReq extends collectionNameReq { transformers?: OutputTransformers; // provide custom data transformer for specific data type like bf16 or f16 vectors } -export interface SearchIteratorReq - extends Omit< - SearchSimpleReq, - 'data' | 'vectors' | 'offset' | 'limit' | 'topk' - > { - data: number[]; // data to search - batchSize: number; - limit: number; -} - export type SearchTextType = string | string[]; export type SearchVectorType = VectorTypes | VectorTypes[]; export type SearchDataType = SearchVectorType | SearchTextType; @@ -338,7 +333,7 @@ export type SearchMultipleDataType = VectorTypes[] | SearchTextType[]; // simplified search api parameter type export interface SearchSimpleReq extends collectionNameReq { partition_names?: string[]; // partition names - anns_field?: string; // your vector field name，rquired if you are searching on multiple vector fields collection + anns_field?: string; // your vector field name，required if you are searching on multiple vector fields collection data?: SearchDataType; // vector or text to search vector?: VectorTypes; // alias for data, deprecated vectors?: VectorTypes[]; // alias for data, deprecated @@ -372,6 +367,12 @@ export type HybridSearchSingleReq = Pick< transformers?: OutputTransformers; // provide custom data transformer for specific data type like bf16 or f16 vectors }; +export interface SearchIteratorReq + extends Omit { + 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 +} + // rerank strategy and parameters export type RerankerObj = { strategy: RANKER_TYPE | string; // rerank strategy @@ -433,7 +434,12 @@ export interface SearchRes extends resStatusResponse { output_fields: string[]; group_by_field_value: string; recalls: number[]; + search_iterator_v2_results?: Record; + _search_iterator_v2_results?: string; + all_search_count?: number; }; + collection_name: string; + session_ts: number; } // because in javascript, there is no float16 and bfloat16 type @@ -462,8 +468,8 @@ export type QueryReq = BaseQueryReq & export interface QueryIteratorReq extends Omit { - limit?: number; - batchSize: number; + 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 } export interface GetReq extends collectionNameReq { diff --git a/milvus/utils/Format.ts b/milvus/utils/Format.ts index 4a0c77e9..fdead122 100644 --- a/milvus/utils/Format.ts +++ b/milvus/utils/Format.ts @@ -681,7 +681,8 @@ export const buildSearchParams = ( ignore_growing: data.ignore_growing ?? false, }; - // if group_by_field is set, add it to the search params + // if group_by_field is set + // reminder: never add this kind of key again, just put params in the params object if (data.group_by_field) { search_params.group_by_field = data.group_by_field; } @@ -695,6 +696,11 @@ export const buildSearchParams = ( search_params.hints = data.hints; } + // data.params -> search_params + for (let key in data.params) { + search_params[key] = data.params[key]; + } + return search_params; }; diff --git a/milvus/utils/Function.ts b/milvus/utils/Function.ts index c3bf4756..be8316a8 100644 --- a/milvus/utils/Function.ts +++ b/milvus/utils/Function.ts @@ -191,27 +191,6 @@ export const getQueryIteratorExpr = (params: { }); }; -// return distance range between the first and last item for the given search results -export const getRangeFromSearchResult = (results: SearchResultData[]) => { - // get first item - const firstItem = results[0]; - const lastItem = results[results.length - 1]; - - if (firstItem && lastItem) { - const radius = lastItem.score * 2 - firstItem.score; - return { - radius: radius, - lastDistance: lastItem.score, - id: lastItem.id, - }; - } else { - return { - radius: 0, - lastDistance: 0, - }; - } -}; - // return pk filed != expression based on pk field type, if pk field is string, return pk field != '' export const getPKFieldExpr = (data: { pkField: FieldSchema; diff --git a/test/grpc/Iterator.spec.ts b/test/grpc/Iterator.spec.ts index f70ccd64..adfa2767 100644 --- a/test/grpc/Iterator.spec.ts +++ b/test/grpc/Iterator.spec.ts @@ -372,210 +372,184 @@ describe(`Iterator API`, () => { }); }); - // it('search iterator with batch size = total should success', async () => { - // const batchSize = 100; - // const total = 100; - // const iterator = await milvusClient.searchIterator({ - // collection_name: COLLECTION, - // batchSize: batchSize, - // data: data[0].vector, - // expr: 'id > 0', - // output_fields: ['id'], - // limit: total, - // }); - - // const results: any = []; - // // let batch = 0; - // for await (const value of iterator) { - // // console.log(`batch${batch++}`, value.length); - // // console.log(value.map((item: any) => item.score)); - // results.push(...value); - // } - - // // results id should be unique - // const idSet = new Set(); - // results.forEach((result: any) => { - // idSet.add(result.id); - // }); - // expect(idSet.size).toEqual(total); - // }); - - // it('search iterator with batch size > total should success', async () => { - // const batchSize = 200; - // const total = 100; - // const iterator = await milvusClient.searchIterator({ - // collection_name: COLLECTION, - // batchSize: batchSize, - // data: data[0].vector, - // expr: 'id > 0', - // output_fields: ['id'], - // limit: total, - // }); - - // const results: any = []; - // // let batch = 0; - // for await (const value of iterator) { - // // console.log(`batch${batch++}`, value.length); - // // console.log(value.map((item: any) => item.score)); - // results.push(...value); - // } - - // // results id should be unique - // const idSet = new Set(); - // results.forEach((result: any) => { - // idSet.add(result.id); - // }); - // expect(idSet.size).toEqual(total); - // }); - - // it('search iterator with batch size < total should success', async () => { - // const batchSize = 33; - // const total = 100; - // const iterator = await milvusClient.searchIterator({ - // collection_name: COLLECTION, - // batchSize: batchSize, - // data: data[0].vector, - // expr: 'id > 0', - // output_fields: ['id'], - // limit: total, - // }); - - // const results: any = []; - // let batchTimes = 0; - // for await (const value of iterator) { - // // console.log(`batch${batch++}`, value.length); - // // console.log(value.map((item: any) => item.score)); - // batchTimes++; - // results.push(...value); - // } - // expect(batchTimes).toEqual(Math.ceil(total / batchSize)); - - // // results id should be unique - // const idSet = new Set(); - // results.forEach((result: any) => { - // idSet.add(result.id); - // }); - // expect(idSet.size).toEqual(total); - // }); - - // it('search iterator with batch size = 2 should success, and ignore total', async () => { - // const batchSize = 2; - // const total = 10; - // const iterator = await milvusClient.searchIterator({ - // collection_name: COLLECTION, - // batchSize: batchSize, - // data: [0.1, 0.2, 0.3, 0.4], - // expr: 'id > 0', - // output_fields: ['id'], - // limit: total, - // }); - - // const results: any = []; - // // let batch = 0; - // for await (const value of iterator) { - // // console.log(`batch${batch++}`, value.length); - // // console.log(value.map((item: any) => item.score)); - // results.push(...value); - // } - - // expect(results.length).toEqual(total); - - // // results id should be unique - // const idSet = new Set(); - // results.forEach((result: any) => { - // idSet.add(result.id); - // }); - // expect(idSet.size).toEqual(total); - // }); - - // it('search iterator with batch size = 1 should success, and ignore total', async () => { - // const search = await milvusClient.search({ - // collection_name: COLLECTION, - // data: [0.1, 0.2, 0.3, 0.4], - // expr: 'id > 0', - // output_fields: ['id'], - // limit: 10, - // }); - - // const batchSize = 1; - // const total = 10; - // const iterator = await milvusClient.searchIterator({ - // collection_name: COLLECTION, - // batchSize: batchSize, - // data: [0.1, 0.2, 0.3, 0.4], - // expr: 'id > 0', - // output_fields: ['id'], - // limit: total, - // }); - - // const results: any = []; - // // let batch = 0; - // for await (const value of iterator) { - // // console.log(`batch${batch++}`, value.length); - // // console.log(value.map((item: any) => item.score)); - // results.push(...value); - // } - - // expect(results.length).toEqual(total); - - // // results id should be unique - // const idSet = new Set(); - // results.forEach((result: any) => { - // idSet.add(result.id); - // }); - // expect(idSet.size).toEqual(total); - - // // get result scores - // const scores = results.map((result: any) => result.score); - - // // compare with search result, should be equal - // expect(scores).toEqual(search.results.map(s => s.score)); - // }); - - // it('search iterator with limit > all data count should success, and ignore total', async () => { - // const batchSize = 5000; - // const total = 30000; - // const iterator = await milvusClient.searchIterator({ - // collection_name: COLLECTION, - // batchSize: batchSize, - // data: data[0].vector, - // expr: 'id > 0', - // output_fields: ['id'], - // limit: total, - // }); - - // const results: any = []; - // // let batch = 0; - // for await (const value of iterator) { - // // console.log(`batch${batch++}`, value.length); - // // console.log(value.map((item: any) => item.score)); - // results.push(...value); - // } - - // expect(results.length).toEqual(data.length); - - // // results id should be unique - // const idSet = new Set(); - // results.forEach((result: any) => { - // idSet.add(result.id); - // }); - // expect(idSet.size).toEqual(data.length); - // }); - - // it('search with cosine similarity should success', async () => { - // const batchSize = 10000; - // const total = 2000; - // const searchRes = await milvusClient.search({ - // collection_name: COLLECTION_COSINE, - // data: cosineData[0].vector, - // expr: 'id > 0', - // output_fields: ['id'], - // limit: total, - // params: { - // radius: 0.6, - // range_filter: 0.5, - // }, - // }); - // console.log(searchRes.results.map(s => s.score)); - // }); + it('search iterator with batch size = 2 should success', async () => { + const batchSize = 2; + const total = 30; + const iterator = await milvusClient.searchIterator({ + collection_name: COLLECTION, + batchSize: batchSize, + data: data[0].vector, + expr: 'id > 0', + output_fields: ['id'], + limit: total, + }); + + const results: any = []; + // let batch = 0; + for await (const value of iterator) { + // console.log(`batch${batch++}`, value.length); + // console.log(value.map((item: any) => item.score)); + results.push(...value); + } + + // // results id should be unique + const idSet = new Set(); + results.forEach((result: any) => { + idSet.add(result.id); + }); + expect(idSet.size).toEqual(total); + }); + + it('search iterator with batch size = total should success', async () => { + const batchSize = 5000; + const total = 5000; + const iterator = await milvusClient.searchIterator({ + collection_name: COLLECTION, + batchSize: batchSize, + data: data[0].vector, + expr: 'id > 0', + output_fields: ['id'], + limit: total, + }); + + const results: any = []; + // let batch = 0; + for await (const value of iterator) { + // console.log(`batch${batch++}`, value.length); + // console.log(value.map((item: any) => item.score)); + results.push(...value); + } + + // // results id should be unique + const idSet = new Set(); + results.forEach((result: any) => { + idSet.add(result.id); + }); + expect(idSet.size).toEqual(total); + }); + + it('search iterator with batch size > total should success', async () => { + const batchSize = 20000; + const total = 10000; + const iterator = await milvusClient.searchIterator({ + collection_name: COLLECTION, + batchSize: batchSize, + data: data[0].vector, + expr: 'id > 0', + output_fields: ['id'], + limit: total, + }); + + const results: any = []; + // let batch = 0; + for await (const value of iterator) { + // console.log(`batch${batch++}`, value.length); + // console.log(value.map((item: any) => item.score)); + results.push(...value); + } + + // // results id should be unique + const idSet = new Set(); + results.forEach((result: any) => { + idSet.add(result.id); + }); + expect(idSet.size).toEqual(total); + }); + + it('search iterator with batch size < total should success', async () => { + const batchSize = 3000; + const total = 10000; + const iterator = await milvusClient.searchIterator({ + collection_name: COLLECTION, + batchSize: batchSize, + data: data[0].vector, + expr: 'id > 0', + output_fields: ['id'], + limit: total, + }); + + const results: any = []; + // let batch = 0; + for await (const value of iterator) { + // console.log(`batch${batch++}`, value.length); + // console.log(value.map((item: any) => item.score)); + results.push(...value); + } + + // // results id should be unique + const idSet = new Set(); + results.forEach((result: any) => { + idSet.add(result.id); + }); + expect(idSet.size).toEqual(total); + }); + + it('search iterator with limit = -1 should success', async () => { + const batchSize = 3000; + const iterator = await milvusClient.searchIterator({ + collection_name: COLLECTION, + batchSize: batchSize, + data: data[0].vector, + expr: 'id > 0', + output_fields: ['id'], + limit: NO_LIMIT, + }); + + const results: any = []; + // let batch = 0; + for await (const value of iterator) { + // console.log(`batch${batch++}`, value.length); + // console.log(value.map((item: any) => item.score)); + results.push(...value); + } + + // // results id should be unique + expect(results.length).toEqual(data.length); + }); + + it('search iterator without limit should success', async () => { + const batchSize = 3000; + const iterator = await milvusClient.searchIterator({ + collection_name: COLLECTION, + batchSize: batchSize, + data: data[0].vector, + expr: 'id > 0', + output_fields: ['id'], + }); + + const results: any = []; + // let batch = 0; + for await (const value of iterator) { + // console.log(`batch${batch++}`, value.length); + // console.log(value.map((item: any) => item.score)); + results.push(...value); + } + + // // results id should be unique + expect(results.length).toEqual(data.length); + }); + + it('search iterator with limit > total should success', async () => { + const batchSize = 3000; + const limit = 30000; + const iterator = await milvusClient.searchIterator({ + collection_name: COLLECTION, + batchSize: batchSize, + data: data[0].vector, + expr: 'id > 0', + output_fields: ['id'], + limit: limit, + }); + + const results: any = []; + // let batch = 0; + for await (const value of iterator) { + // console.log(`batch${batch++}`, value.length); + // console.log(value.map((item: any) => item.score)); + results.push(...value); + } + + // // results id should be unique + expect(results.length).toEqual(data.length); + }); }); diff --git a/test/utils/Format.spec.ts b/test/utils/Format.spec.ts index 23395b8f..2564c063 100644 --- a/test/utils/Format.spec.ts +++ b/test/utils/Format.spec.ts @@ -34,8 +34,6 @@ import { SearchSimpleReq, formatExprValues, } from '../../milvus'; -import { json } from 'stream/consumers'; -import exp from 'constants'; describe('utils/format', () => { it(`all kinds of url should be supported`, async () => { @@ -1058,13 +1056,14 @@ describe('utils/format', () => { offset: 0, metric_type: '', ignore_growing: false, + nprobe: 2, }); const data2: SearchSimpleReq = { collection_name: 'test', data: [1, 2, 3, 4, 5, 6, 7, 8], anns_field: 'vector', - params: { nprobe: 2 }, + params: { nprobe: 2, test: 'test' }, limit: 2, output_fields: ['vector', 'vector1'], group_by_field: 'group_by_field_value', @@ -1073,10 +1072,10 @@ describe('utils/format', () => { }; const newSearchParams2 = buildSearchParams(data2, anns_field); - + console.dir(newSearchParams2, { depth: null }); expect(newSearchParams2).toEqual({ anns_field: 'vector', - params: '{"nprobe":2}', + params: '{"nprobe":2,"test":"test"}', topk: 2, offset: 0, metric_type: '', @@ -1084,6 +1083,8 @@ describe('utils/format', () => { group_by_field: 'group_by_field_value', group_size: 5, strict_group_size: true, + test: 'test', + nprobe: 2, }); }); diff --git a/test/utils/Function.spec.ts b/test/utils/Function.spec.ts index 035fca5d..ce55e4a1 100644 --- a/test/utils/Function.spec.ts +++ b/test/utils/Function.spec.ts @@ -4,7 +4,6 @@ import { DataTypeStringEnum, DEFAULT_MIN_INT64, getPKFieldExpr, - getRangeFromSearchResult, SearchResultData, getSparseDim, SparseFloatVector, @@ -151,138 +150,6 @@ describe('Function API testing', () => { expect(result).toBe('id > 10 && field > 10'); }); - it('should return 0 radius when results are empty', () => { - const results = [] as any; - - const result = getRangeFromSearchResult(results); - - expect(result).toEqual({ - radius: 0, - lastDistance: 0, - }); - }); - - it('should return radius and lastDistance when results are not empty', () => { - const results: SearchResultData[] = [ - { - id: '1', - score: 0.1, - }, - { - id: '2', - score: 0.2, - }, - { - id: '3', - score: 0.3, - }, - ]; - - const result = getRangeFromSearchResult(results); - - expect(result).toEqual({ - radius: 0.3 * 2 - 0.1, - lastDistance: 0.3, - id: '3', - }); - }); - - it('should return 0 radius when results contain only one item', () => { - const results: SearchResultData[] = [ - { - id: '1', - score: 0.1, - }, - ]; - - const result = getRangeFromSearchResult(results); - - expect(result).toEqual({ - radius: 0.1 * 2 - 0.1, - lastDistance: 0.1, - id: '1', - }); - }); - - it('should return 0 radius when results contain only two items', () => { - const results: SearchResultData[] = [ - { - id: '1', - score: 0.1, - }, - { - id: '2', - score: 0.2, - }, - ]; - - const result = getRangeFromSearchResult(results); - - expect(result).toEqual({ - radius: 0.2 * 2 - 0.1, - lastDistance: 0.2, - id: '2', - }); - }); - - it('should return varchar expression when pk field is varchar', () => { - const pkField: any = { - name: 'id', - data_type: DataTypeStringEnum.VarChar, - }; - - const result = getPKFieldExpr({ - pkField, - value: 'abc', - }); - - expect(result).toBe("id != 'abc'"); - }); - - it('should return int64 expression when pk field is int64', () => { - const pkField: any = { - name: 'id', - data_type: DataTypeStringEnum.Int64, - }; - - const result = getPKFieldExpr({ - pkField, - value: 10, - }); - - expect(result).toBe('id != 10'); - }); - - it('should return int64 expression with condition when condition is provided', () => { - const pkField: any = { - name: 'id', - data_type: DataTypeStringEnum.Int64, - }; - - const result = getPKFieldExpr({ - pkField, - value: 10, - condition: '>', - }); - - expect(result).toBe('id > 10'); - }); - - it('should return int64 expression with condition and expr when expr is provided', () => { - const pkField: any = { - name: 'id', - data_type: DataTypeStringEnum.Int64, - }; - - const result = getPKFieldExpr({ - pkField, - value: 10, - condition: '>', - expr: 'field > 10', - }); - - expect(result).toBe('id > 10 && field > 10'); - }); it('should return the correct dimension of the sparse vector', () => { const data = [ { '0': 1, '1': 2, '2': 3 }, From 2f92de2acc1a7e249176fe3aca5bb99b12e131d0 Mon Sep 17 00:00:00 2001 From: shanghaikid Date: Fri, 10 Jan 2025 14:20:01 +0800 Subject: [PATCH 2/4] upgrade milvus test version Signed-off-by: shanghaikid --- package.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/package.json b/package.json index c0c16a0b..3dae9126 100644 --- a/package.json +++ b/package.json @@ -1,7 +1,7 @@ { "name": "@zilliz/milvus2-sdk-node", "author": "Zilliz", - "milvusVersion": "v2.5.1", + "milvusVersion": "v2.5.2", "version": "2.5.3", "main": "dist/milvus", "files": [ From 8f6ac13fd8a5d8b08b5dcd62618c98e9a37bf9e1 Mon Sep 17 00:00:00 2001 From: shanghaikid Date: Fri, 10 Jan 2025 15:44:52 +0800 Subject: [PATCH 3/4] fix test Signed-off-by: shanghaikid --- milvus/grpc/Data.ts | 4 ++-- test/grpc/Iterator.spec.ts | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/milvus/grpc/Data.ts b/milvus/grpc/Data.ts index 6bf1838b..401b29a4 100644 --- a/milvus/grpc/Data.ts +++ b/milvus/grpc/Data.ts @@ -621,8 +621,8 @@ export class Data extends Collection { params[ITER_SEARCH_ID_KEY] = 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 || 0; + batchRes.search_iterator_v2_results?.last_bound; + params[GUARANTEE_TIMESTAMP_KEY] = batchRes.session_ts; params[ITER_SEARCH_BATCH_SIZE_KEY] = batchSize; return { diff --git a/test/grpc/Iterator.spec.ts b/test/grpc/Iterator.spec.ts index adfa2767..e58bc9c4 100644 --- a/test/grpc/Iterator.spec.ts +++ b/test/grpc/Iterator.spec.ts @@ -372,8 +372,8 @@ describe(`Iterator API`, () => { }); }); - it('search iterator with batch size = 2 should success', async () => { - const batchSize = 2; + it('search iterator with batch size = 1 should success', async () => { + const batchSize = 1; const total = 30; const iterator = await milvusClient.searchIterator({ collection_name: COLLECTION, From a448b93e32b3525df6efaeb00eeba29134348a27 Mon Sep 17 00:00:00 2001 From: shanghaikid Date: Fri, 10 Jan 2025 16:02:43 +0800 Subject: [PATCH 4/4] change log level Signed-off-by: shanghaikid --- test/grpc/Basic.spec.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/grpc/Basic.spec.ts b/test/grpc/Basic.spec.ts index 13e803f0..a78caaba 100644 --- a/test/grpc/Basic.spec.ts +++ b/test/grpc/Basic.spec.ts @@ -3,7 +3,7 @@ import { IP, GENERATE_NAME, generateInsertData } from '../tools'; const milvusClient = new MilvusClient({ address: IP, - logLevel: 'debug', + logLevel: 'info', logPrefix: 'Basic API', }); const COLLECTION_NAME = GENERATE_NAME();