Skip to content

Commit

Permalink
feat: support partial load (#355)
Browse files Browse the repository at this point in the history
* support partial load fields

Signed-off-by: ryjiang <jiangruiyi@gmail.com>

* finish partial load

Signed-off-by: ryjiang <jiangruiyi@gmail.com>

---------

Signed-off-by: ryjiang <jiangruiyi@gmail.com>
  • Loading branch information
shanghaikid authored Sep 2, 2024
1 parent 3448094 commit 2b9ba0b
Show file tree
Hide file tree
Showing 4 changed files with 187 additions and 60 deletions.
2 changes: 2 additions & 0 deletions milvus/types/Collection.ts
Original file line number Diff line number Diff line change
Expand Up @@ -127,6 +127,8 @@ export interface LoadCollectionReq extends collectionNameReq {
replica_number?: number; // optional, replica number, default is 1
resource_groups?: string[]; // optional, resource groups
refresh?: boolean; // optional, refresh, default is false
load_fields?: string[]; // optional, load fields
skip_load_dynamic_field?: boolean; // optional, skip load dynamic field, default is false
}
export interface ReleaseLoadCollectionReq extends collectionNameReq {}

Expand Down
109 changes: 50 additions & 59 deletions milvus/utils/Format.ts
Original file line number Diff line number Diff line change
Expand Up @@ -693,71 +693,62 @@ export const buildSearchRequest = (
searchHybridReq.data[0].anns_field
);

// output fields(reference fields)
const default_output_fields: string[] = [];

// Iterate through collection fields, create search request
for (let i = 0; i < collectionInfo.schema.fields.length; i++) {
const field = collectionInfo.schema.fields[i];
const { name, dataType } = field;

// if field type is vector, build the request
if (isVectorType(dataType)) {
let req: SearchSimpleReq | (HybridSearchReq & HybridSearchSingleReq) =
data as SearchSimpleReq;

if (isHybridSearch) {
const singleReq = searchHybridReq.data.find(d => d.anns_field === name);
// if it is hybrid search and no request target is not found, skip
if (!singleReq) {
continue;
}
// merge single request with hybrid request
req = Object.assign(cloneObj(data), singleReq);
} else {
// if it is not hybrid search, and we have built one request, skip
const skip =
requests.length === 1 ||
(typeof req.anns_field !== 'undefined' && req.anns_field !== name);
if (skip) {
continue;
}
}
const { name } = field;

// get search vectors
let searchingVector: VectorTypes | VectorTypes[] = isHybridSearch
? req.data!
: searchReq.vectors ||
searchSimpleReq.vectors ||
searchSimpleReq.vector ||
searchSimpleReq.data;

// format searching vector
searchingVector = formatSearchVector(searchingVector, field.dataType!);

// create search request
requests.push({
collection_name: req.collection_name,
partition_names: req.partition_names || [],
output_fields: req.output_fields || default_output_fields,
nq: searchReq.nq || searchingVector.length,
dsl: searchReq.expr || searchSimpleReq.filter || '',
dsl_type: DslType.BoolExprV1,
placeholder_group: buildPlaceholderGroupBytes(
milvusProto,
searchingVector as VectorTypes[],
field.dataType!
),
search_params: parseToKeyValue(
searchReq.search_params || buildSearchParams(req, name)
),
consistency_level:
req.consistency_level || (collectionInfo.consistency_level as any),
});
let req: SearchSimpleReq | (HybridSearchReq & HybridSearchSingleReq) =
data as SearchSimpleReq;

if (isHybridSearch) {
const singleReq = searchHybridReq.data.find(d => d.anns_field === name);
// if it is hybrid search and no request target is not found, skip
if (!singleReq) {
continue;
}
// merge single request with hybrid request
req = Object.assign(cloneObj(data), singleReq);
} else {
// if field is not vector, add it to output fields
default_output_fields.push(name);
// if it is not hybrid search, and we have built one request, skip
const skip =
requests.length === 1 ||
(typeof req.anns_field !== 'undefined' && req.anns_field !== name);
if (skip) {
continue;
}
}

// get search vectors
let searchingVector: VectorTypes | VectorTypes[] = isHybridSearch
? req.data!
: searchReq.vectors ||
searchSimpleReq.vectors ||
searchSimpleReq.vector ||
searchSimpleReq.data;

// format searching vector
searchingVector = formatSearchVector(searchingVector, field.dataType!);

// create search request
requests.push({
collection_name: req.collection_name,
partition_names: req.partition_names || [],
output_fields: req.output_fields || ['*'],
nq: searchReq.nq || searchingVector.length,
dsl: searchReq.expr || searchSimpleReq.filter || '',
dsl_type: DslType.BoolExprV1,
placeholder_group: buildPlaceholderGroupBytes(
milvusProto,
searchingVector as VectorTypes[],
field.dataType!
),
search_params: parseToKeyValue(
searchReq.search_params || buildSearchParams(req, name)
),
consistency_level:
req.consistency_level || (collectionInfo.consistency_level as any),
});
}

/**
Expand Down
2 changes: 1 addition & 1 deletion package.json
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
"name": "@zilliz/milvus2-sdk-node",
"author": "ued@zilliz.com",
"version": "2.4.6",
"milvusVersion": "v2.4.8",
"milvusVersion": "v2.4.10",
"main": "dist/milvus",
"files": [
"dist"
Expand Down
134 changes: 134 additions & 0 deletions test/grpc/PartialLoad.spec.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,134 @@
import { MilvusClient, ErrorCode, DataType } from '../../milvus';
import { IP, GENERATE_NAME, generateInsertData } from '../tools';

const milvusClient = new MilvusClient({ address: IP, logLevel: 'info' });
const COLLECTION_NAME = GENERATE_NAME();
const schema = [
{
name: 'vector',
description: 'Vector field',
data_type: DataType.FloatVector,
dim: Number(4),
},
{
name: 'id',
description: 'ID field',
data_type: DataType.Int64,
is_primary_key: true,
autoID: true,
},
{
name: 'varChar',
description: 'VarChar field',
data_type: DataType.VarChar,
max_length: 128,
is_partition_key: false,
},
{
name: 'array',
description: 'array field',
data_type: DataType.Array,
element_type: DataType.VarChar,
max_capacity: 128,
max_length: 128,
is_partition_key: false,
},
];

describe(`Partial load API `, () => {
it(`Create collection should be successful`, async () => {
const res = await milvusClient.createCollection({
collection_name: COLLECTION_NAME,
fields: schema,
});
expect(res.error_code).toEqual(ErrorCode.SUCCESS);
});

it(`Describe collection should be successful`, async () => {
const desc = await milvusClient.describeCollection({
collection_name: COLLECTION_NAME,
});
expect(desc.schema.fields.length).toEqual(schema.length);
});

it(`Create index should be successful`, async () => {
const index = await milvusClient.createIndex({
collection_name: COLLECTION_NAME,
field_name: 'vector',
extra_params: {
index_type: 'IVF_FLAT',
metric_type: 'L2',
params: JSON.stringify({ nlist: 1024 }),
},
});
expect(index.error_code).toEqual(ErrorCode.SUCCESS);
});

it(`load collection should be successful`, async () => {
const load = await milvusClient.loadCollectionSync({
collection_name: COLLECTION_NAME,
load_fields: ['vector', 'id', 'varChar'],
});
expect(load.error_code).toEqual(ErrorCode.SUCCESS);
});

it(`insert data should be successful`, async () => {
const data = generateInsertData(schema);
const insert = await milvusClient.insert({
collection_name: COLLECTION_NAME,
data,
});
expect(insert.status.error_code).toEqual(ErrorCode.SUCCESS);
});

it(`query should be successful`, async () => {
const query = await milvusClient.query({
collection_name: COLLECTION_NAME,
expr: `id > 0`,
output_fields: ['*'],
});
console.dir(query, { depth: null });
expect(query.status.error_code).toEqual(ErrorCode.SUCCESS);
// result should not contain 'array' field
const keys = Object.keys(query.data[0]);
expect(keys.length).toEqual(3);
expect(keys.includes('array')).toBeFalsy();
});

it(`search should be successful`, async () => {
const search = await milvusClient.search({
collection_name: COLLECTION_NAME,
data: [1, 2, 3, 4],
});
expect(search.status.error_code).toEqual(ErrorCode.SUCCESS);
});

it(`search nq > 1 should be successful`, async () => {
const search = await milvusClient.search({
collection_name: COLLECTION_NAME,
data: [
[1, 2, 3, 4],
[5, 6, 7, 8],
],
});
expect(search.status.error_code).toEqual(ErrorCode.SUCCESS);
expect(search.results.length).toEqual(2);
expect(search.results[0].length).toEqual(10);
expect(search.results[1].length).toEqual(10);
});

it(`release and drop should be successful`, async () => {
// releases
const release = await milvusClient.releaseCollection({
collection_name: COLLECTION_NAME,
timeout: 15000,
});
expect(release.error_code).toEqual(ErrorCode.SUCCESS);

// drop
const drop = await milvusClient.dropCollection({
collection_name: COLLECTION_NAME,
});
expect(drop.error_code).toEqual(ErrorCode.SUCCESS);
});
});

0 comments on commit 2b9ba0b

Please sign in to comment.