Skip to content

Commit

Permalink
Support expression template values
Browse files Browse the repository at this point in the history
Signed-off-by: yhmo <[email protected]>
  • Loading branch information
yhmo authored and xiaocai2333 committed Oct 28, 2024
1 parent 3d6de38 commit b97213a
Show file tree
Hide file tree
Showing 7 changed files with 140 additions and 11 deletions.
14 changes: 10 additions & 4 deletions src/main/java/io/milvus/v2/service/vector/VectorService.java
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

public class VectorService extends BaseService {
Expand Down Expand Up @@ -224,12 +225,17 @@ public DeleteResp delete(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStu
if (request.getFilter() == null) {
request.setFilter(vectorUtils.getExprById(respR.getPrimaryFieldName(), request.getIds()));
}
DeleteRequest deleteRequest = DeleteRequest.newBuilder()
DeleteRequest.Builder builder = DeleteRequest.newBuilder()
.setCollectionName(request.getCollectionName())
.setPartitionName(request.getPartitionName())
.setExpr(request.getFilter())
.build();
MutationResult response = blockingStub.delete(deleteRequest);
.setExpr(request.getFilter());
if (request.getFilter() != null && !request.getFilter().isEmpty()) {
Map<String, Object> filterTemplateValues = request.getFilterTemplateValues();
filterTemplateValues.forEach((key, value)->{
builder.putExprTemplateValues(key, vectorUtils.deduceAndCreateTemplateValue(value));
});
}
MutationResult response = blockingStub.delete(builder.build());
rpcUtils.handleResponse(title, response.getStatus());
GTsDict.getInstance().updateCollectionTs(request.getCollectionName(), response.getTimestamp());
return DeleteResp.builder()
Expand Down
13 changes: 13 additions & 0 deletions src/main/java/io/milvus/v2/service/vector/request/DeleteReq.java
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,9 @@
import lombok.Data;
import lombok.experimental.SuperBuilder;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

@Data
@SuperBuilder
Expand All @@ -33,4 +35,15 @@ public class DeleteReq {
private String partitionName = "";
private String filter;
private List<Object> ids;

// Expression template, to improve expression parsing performance in complicated list
// Assume user has a filter = "pk > 3 and city in ["beijing", "shanghai", ......]
// The long list of city will increase the time cost to parse this expression.
// So, we provide exprTemplateValues for this purpose, user can set filter like this:
// filter = "pk > {age} and city in {city}"
// filterTemplateValues = Map{"age": 3, "city": List<String>{"beijing", "shanghai", ......}}
// Valid value of this map can be:
// Boolean, Long, Double, String, List<Boolean>, List<Long>, List<Double>, List<String>
@Builder.Default
private Map<String, Object> filterTemplateValues = new HashMap<>();
}
15 changes: 12 additions & 3 deletions src/main/java/io/milvus/v2/service/vector/request/QueryReq.java
Original file line number Diff line number Diff line change
Expand Up @@ -24,9 +24,7 @@
import lombok.Data;
import lombok.experimental.SuperBuilder;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.*;

@Data
@SuperBuilder
Expand All @@ -42,4 +40,15 @@ public class QueryReq {
private ConsistencyLevel consistencyLevel = null;
private long offset;
private long limit;

// Expression template, to improve expression parsing performance in complicated list
// Assume user has a filter = "pk > 3 and city in ["beijing", "shanghai", ......]
// The long list of city will increase the time cost to parse this expression.
// So, we provide exprTemplateValues for this purpose, user can set filter like this:
// filter = "pk > {age} and city in {city}"
// filterTemplateValues = Map{"age": 3, "city": List<String>{"beijing", "shanghai", ......}}
// Valid value of this map can be:
// Boolean, Long, Double, String, List<Boolean>, List<Long>, List<Double>, List<String>
@Builder.Default
private Map<String, Object> filterTemplateValues = new HashMap<>();
}
11 changes: 11 additions & 0 deletions src/main/java/io/milvus/v2/service/vector/request/SearchReq.java
Original file line number Diff line number Diff line change
Expand Up @@ -57,4 +57,15 @@ public class SearchReq {
private ConsistencyLevel consistencyLevel = null;
private boolean ignoreGrowing;
private String groupByFieldName;

// Expression template, to improve expression parsing performance in complicated list
// Assume user has a filter = "pk > 3 and city in ["beijing", "shanghai", ......]
// The long list of city will increase the time cost to parse this expression.
// So, we provide exprTemplateValues for this purpose, user can set filter like this:
// filter = "pk > {age} and city in {city}"
// filterTemplateValues = Map{"age": 3, "city": List<String>{"beijing", "shanghai", ......}}
// Valid value of this map can be:
// Boolean, Long, Double, String, List<Boolean>, List<Long>, List<Double>, List<String>
@Builder.Default
private Map<String, Object> filterTemplateValues = new HashMap<>();
}
61 changes: 61 additions & 0 deletions src/main/java/io/milvus/v2/utils/VectorUtils.java
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,12 @@ public QueryRequest ConvertToGrpcQueryRequest(QueryReq request){
.addAllPartitionNames(request.getPartitionNames())
.addAllOutputFields(request.getOutputFields())
.setExpr(request.getFilter());
if (request.getFilter() != null && !request.getFilter().isEmpty()) {
Map<String, Object> filterTemplateValues = request.getFilterTemplateValues();
filterTemplateValues.forEach((key, value)->{
builder.putExprTemplateValues(key, deduceAndCreateTemplateValue(value));
});
}

// a new parameter from v2.2.9, if user didn't specify consistency level, set this parameter to true
if (request.getConsistencyLevel() == null) {
Expand Down Expand Up @@ -180,6 +186,10 @@ public SearchRequest ConvertToGrpcSearchRequest(SearchReq request) {
builder.setDslType(DslType.BoolExprV1);
if (request.getFilter() != null && !request.getFilter().isEmpty()) {
builder.setDsl(request.getFilter());
Map<String, Object> filterTemplateValues = request.getFilterTemplateValues();
filterTemplateValues.forEach((key, value)->{
builder.putExprTemplateValues(key, deduceAndCreateTemplateValue(value));
});
}

long guaranteeTimestamp = getGuaranteeTimestamp(request.getConsistencyLevel(), request.getCollectionName());
Expand All @@ -195,6 +205,57 @@ public SearchRequest ConvertToGrpcSearchRequest(SearchReq request) {
return builder.build();
}

public static TemplateValue deduceAndCreateTemplateValue(Object value) {
if (value instanceof Boolean) {
return TemplateValue.newBuilder()
.setBoolVal((Boolean)value)
.setType(DataType.Bool)
.build();
} else if (value instanceof Integer) {
return TemplateValue.newBuilder()
.setInt64Val((Integer)value)
.setType(DataType.Int64)
.build();
} else if (value instanceof Double) {
return TemplateValue.newBuilder()
.setFloatVal((Double)value)
.setType(DataType.Double)
.build();
} else if (value instanceof String) {
return TemplateValue.newBuilder()
.setStringVal((String)value)
.setType(DataType.VarChar)
.build();
} else if (value instanceof List) {
// TemplateArrayValue and TemplateValue can nest each other
// The element_type of TemplateArrayValue is deduced by its elements:
// 1. if all the elements are the same type, element_type is the first element's type
// 2. if not the same type, element_type is DataType.JSON
List<Object> array = (List<Object>)value;
TemplateArrayValue.Builder builder = TemplateArrayValue.newBuilder();
DataType lastType = DataType.UNRECOGNIZED;
boolean sameType = true;
for (Object obj : array) {
TemplateValue tv = deduceAndCreateTemplateValue(obj);
builder.addArray(tv);
if (sameType && lastType != DataType.UNRECOGNIZED && lastType != tv.getType()) {
sameType = false;
}
lastType = tv.getType();
}
DataType arrayType = sameType ? lastType : DataType.JSON;
builder.setElementType(arrayType);
builder.setSameType(sameType);

return TemplateValue.newBuilder()
.setArrayVal(builder.build())
.setType(builder.getElementType())
.build();
} else {
throw new ParamException("Unsupported value type for expression template.");
}
}

public static SearchRequest convertAnnSearchParam(@NonNull AnnSearchReq annSearchReq,
ConsistencyLevel consistencyLevel) {
SearchRequest.Builder builder = SearchRequest.newBuilder();
Expand Down
2 changes: 1 addition & 1 deletion src/main/milvus-proto
35 changes: 32 additions & 3 deletions src/test/java/io/milvus/v2/service/vector/VectorTest.java
Original file line number Diff line number Diff line change
Expand Up @@ -29,9 +29,7 @@
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.*;

class VectorTest extends BaseTest {

Expand Down Expand Up @@ -105,6 +103,37 @@ void testSearch() {
logger.info(statusR.toString());
}

@Test
void testSearchWithTemplateExpression() {
List<Float> vectorList = new ArrayList<>();
vectorList.add(1.0f);
vectorList.add(2.0f);

Map<String, Map<String, Object>> expressionTemplateValues = new HashMap<>();
Map<String, Object> params = new HashMap<>();
params.put("max", 10);
expressionTemplateValues.put("id < {max}", params);

List<Object> list = Arrays.asList(1, 2, 3);
Map<String, Object> params2 = new HashMap<>();
params2.put("list", list);
expressionTemplateValues.put("id in {list}", params2);

expressionTemplateValues.forEach((key, value) -> {
SearchReq request = SearchReq.builder()
.collectionName("test")
.data(Collections.singletonList(new FloatVec(vectorList)))
.topK(10)
.offset(0L)
.filter(key)
.filterTemplateValues(value)
.build();
SearchResp statusR = client_v2.search(request);
logger.info(statusR.toString());
System.out.println(statusR.toString());
});
}

@Test
void testDelete() {
DeleteReq request = DeleteReq.builder()
Expand Down

0 comments on commit b97213a

Please sign in to comment.