Skip to content

Commit

Permalink
[2.4] support expr template values for search/query/delete (#391)
Browse files Browse the repository at this point in the history 
* add formatExprValues formatter

Signed-off-by: ryjiang <[email protected]>

* finish

Signed-off-by: ryjiang <[email protected]>

---------

Signed-off-by: ryjiang <[email protected]>
  • Loading branch information
@shanghaikid
shanghaikid authored Dec 23, 2024
1 parent 646418e commit f593f90
Show file tree
Hide file tree
Showing 9 changed files with 451 additions and 47 deletions.
23 changes: 19 additions & 4 deletions milvus/grpc/Data.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,7 @@ import {
f32ArrayToBinaryBytes,
NO_LIMIT,
DescribeCollectionReq,
formatExprValues,
} from '../';
import { Collection } from './Collection';

Expand Down Expand Up @@ -345,10 +346,17 @@ export class Data extends Collection {
// filter > expr
data.expr = data.filter || data.expr;

const req = data as any;

// if exprValues exist, format it
if (data.exprValues) {
req.expr_template_values = formatExprValues(data.exprValues);
}

const promise = await promisify(
this.channelPool,
'Delete',
data,
req,
data.timeout || this.timeout
);
return promise;
Expand Down Expand Up @@ -410,7 +418,7 @@ export class Data extends Collection {
if ((data as DeleteByFilterReq).filter) {
expr = (data as DeleteByFilterReq).filter;
}
const req = { ...data, expr };
const req = { ...data, expr } as any;
return this.deleteEntities(req);
}

Expand Down Expand Up @@ -459,7 +467,9 @@ export class Data extends Collection {
describeCollectionRequest.db_name = data.db_name;
}

const collectionInfo = await this.describeCollection(describeCollectionRequest);
const collectionInfo = await this.describeCollection(
describeCollectionRequest
);

// build search params
const { request, nq, round_decimal, isHybridSearch } = buildSearchRequest(
Expand Down Expand Up @@ -755,7 +765,7 @@ export class Data extends Collection {
});

// search data
const res = await client.query(data);
const res = await client.query(data as QueryReq);

// get last item of the data
const lastItem = res.data[res.data.length - 1];
Expand Down Expand Up @@ -926,6 +936,11 @@ export class Data extends Collection {
// filter > expr or empty > ids
data.expr = data.filter || data.expr || primaryKeyInIdsExpression;

// if exprValues exist, format it
if (data.exprValues) {
(data as any).expr_template_values = formatExprValues(data.exprValues);
}

// Execute the query and get the results
const promise: QueryRes = await promisify(
this.channelPool,
Expand Down
5 changes: 4 additions & 1 deletion milvus/types/Common.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,10 @@ export interface TimeStampArray {
created_utc_timestamps: string[];
}

export type keyValueObj = Record<string, string | number | string[] | number[]>;
export type keyValueObj = Record<
string,
string | number | string[] | number[] | unknown
>;

export interface collectionNameReq extends GrpcTimeOut {
collection_name: string; // required, collection name
Expand Down
20 changes: 13 additions & 7 deletions milvus/types/Data.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -120,12 +120,11 @@ interface BaseDeleteReq extends collectionNameReq {
| 'Bounded'
| 'Eventually'
| 'Customized'; // consistency level
exprValues?: keyValueObj; // template values for filter expression, eg: {key: 'value'}
}

export interface DeleteEntitiesReq extends BaseDeleteReq {
filter?: string; // filter expression
expr?: string; // alias for filter
}
export type DeleteEntitiesReq = BaseDeleteReq &
({ expr?: string; filter?: never } | { filter?: string; expr?: never });

export interface DeleteByIdsReq extends BaseDeleteReq {
ids: string[] | number[]; // primary key values
Expand Down Expand Up @@ -288,6 +287,7 @@ export interface SearchReq extends collectionNameReq {
anns_field?: string; // your vector field name
partition_names?: string[]; // partition names
expr?: string; // filter expression
exprValues?: keyValueObj; // template values for filter expression, eg: {key: 'value'}
search_params: SearchParam; // search parameters
vectors: VectorTypes[]; // vectors to search
output_fields?: string[]; // fields to return
Expand Down Expand Up @@ -321,6 +321,7 @@ export interface SearchSimpleReq extends collectionNameReq {
offset?: number; // skip how many results
filter?: string; // filter expression
expr?: string; // alias for filter
exprValues?: keyValueObj; // template values for filter expression, eg: {key: 'value'}
params?: keyValueObj; // extra search parameters
metric_type?: string; // distance metric type
consistency_level?: ConsistencyLevelEnum; // consistency level
Expand All @@ -336,6 +337,7 @@ export type HybridSearchSingleReq = Pick<
> & {
data: VectorTypes[] | VectorTypes; // vector to search
expr?: string; // filter expression
exprValues?: keyValueObj; // template values for filter expression, eg: {key: 'value'}
params?: keyValueObj; // extra search parameters
transformers?: OutputTransformers; // provide custom data transformer for specific data type like bf16 or f16 vectors
};
Expand Down Expand Up @@ -406,17 +408,21 @@ export type OutputTransformers = {
[DataType.SparseFloatVector]?: (sparse: SparseVectorDic) => SparseFloatVector;
};

export interface QueryReq extends collectionNameReq {
type BaseQueryReq = collectionNameReq & {
output_fields?: string[]; // fields to return
partition_names?: string[]; // partition names
ids?: string[] | number[]; // primary key values
expr?: string; // filter expression
expr?: string; // filter expression, or template string, eg: "key = {key}"
filter?: string; // alias for expr
offset?: number; // skip how many results
limit?: number; // how many results you want
consistency_level?: ConsistencyLevelEnum; // consistency level
transformers?: OutputTransformers; // provide custom data transformer for specific data type like bf16 or f16 vectors
}
exprValues?: keyValueObj; // template values for filter expression, eg: {key: 'value'}
};

export type QueryReq = BaseQueryReq &
({ expr?: string; filter?: never } | { filter?: string; expr?: never });

export interface QueryIteratorReq
extends Omit<QueryReq, 'ids' | 'offset' | 'limit'> {
Expand Down
134 changes: 118 additions & 16 deletions milvus/utils/Format.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@ import {
f32ArrayToF16Bytes,
bf16BytesToF32Array,
f16BytesToF32Array,
keyValueObj,
} from '../';

/**
Expand Down Expand Up @@ -643,6 +644,21 @@ export const convertRerankParams = (rerank: RerankerObj) => {
return r;
};

type FormatedSearchRequest = {
collection_name: string;
partition_names: string[];
output_fields: string[];
nq?: number;
dsl?: string;
dsl_type?: DslType;
placeholder_group?: Uint8Array;
search_params?: KeyValuePair[];
consistency_level: ConsistencyLevelEnum;
expr?: string;
expr_template_values?: keyValueObj;
rank_params?: KeyValuePair[];
};

/**
* This method is used to build search request for a given data.
* It first fetches the collection info and then constructs the search request based on the data type.
Expand Down Expand Up @@ -676,17 +692,7 @@ export const buildSearchRequest = (
const searchSimpleReq = data as SearchSimpleReq;

// Initialize requests array
const requests: {
collection_name: string;
partition_names: string[];
output_fields: string[];
nq: number;
dsl: string;
dsl_type: DslType;
placeholder_group: Uint8Array;
search_params: KeyValuePair[];
consistency_level: ConsistencyLevelEnum;
}[] = [];
const requests: FormatedSearchRequest[] = [];

// detect if the request is hybrid search request
const isHybridSearch = !!(
Expand Down Expand Up @@ -739,12 +745,12 @@ export const buildSearchRequest = (
searchingVector = formatSearchVector(searchingVector, field.dataType!);

// create search request
requests.push({
const request: FormatedSearchRequest = {
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: req.expr || searchReq.expr || searchSimpleReq.filter || '', // expr
dsl_type: DslType.BoolExprV1,
placeholder_group: buildPlaceholderGroupBytes(
milvusProto,
Expand All @@ -756,7 +762,14 @@ export const buildSearchRequest = (
),
consistency_level:
req.consistency_level || (collectionInfo.consistency_level as any),
});
};

// if exprValues is set, add it to the request(inner)
if (req.exprValues) {
request.expr_template_values = formatExprValues(req.exprValues);
}

requests.push(request);
}
}

Expand All @@ -774,7 +787,7 @@ export const buildSearchRequest = (
return {
isHybridSearch,
request: isHybridSearch
? {
? ({
collection_name: data.collection_name,
partition_names: data.partition_names,
requests: requests,
Expand All @@ -791,7 +804,7 @@ export const buildSearchRequest = (
],
output_fields: requests[0]?.output_fields,
consistency_level: requests[0]?.consistency_level,
}
} as FormatedSearchRequest)
: requests[0],
nq: requests[0].nq,
round_decimal,
Expand Down Expand Up @@ -910,3 +923,92 @@ export const formatSearchVector = (
return searchVector as VectorTypes[];
}
};

type TemplateValue =
| { bool_val: boolean }
| { int64_val: number }
| { float_val: number }
| { string_val: string }
| { array_val: TemplateArrayValue };
type TemplateArrayValue =
| { bool_data: { data: boolean[] } }
| { long_data: { data: number[] } }
| { double_data: { data: number[] } }
| { string_data: { data: string[] } }
| { json_data: { data: any[] } }
| { array_data: { data: TemplateArrayValue[] } };

export const formatExprValues = (
exprValues: Record<string, any>
): Record<string, TemplateValue> => {
const result: Record<string, TemplateValue> = {};
for (const [key, value] of Object.entries(exprValues)) {
if (Array.isArray(value)) {
// Handle arrays
result[key] = { array_val: convertArray(value) };
} else {
// Handle primitive types
if (typeof value === 'boolean') {
result[key] = { bool_val: value };
} else if (typeof value === 'number') {
result[key] = Number.isInteger(value)
? { int64_val: value }
: { float_val: value };
} else if (typeof value === 'string') {
result[key] = { string_val: value };
}
}
}
return result;
};
const convertArray = (arr: any[]): TemplateArrayValue => {
const first = arr[0];
switch (typeof first) {
case 'boolean':
return {
bool_data: {
data: arr,
},
};
case 'number':
if (Number.isInteger(first)) {
return {
long_data: {
data: arr,
},
};
} else {
return {
double_data: {
data: arr,
},
};
}
case 'string':
return {
string_data: {
data: arr,
},
};
case 'object':
if (Array.isArray(first)) {
return {
array_data: {
data: arr.map(convertArray),
},
};
} else {
return {
json_data: {
data: arr,
},
};
}
default:
return {
string_data: {
data: arr,
},
};
}
};
Loading

0 comments on commit f593f90

Please sign in to comment.