Skip to content

Commit 34bda98

Browse files
yhmoyhmo
authored
Add timezone/filterTemplate for QueryIterator/SearchIterator (#1692)
Signed-off-by: yhmo <yihua.mo@zilliz.com>
1 parent c3d162f commit 34bda98

11 files changed

Lines changed: 301 additions & 35 deletions

File tree

‎examples/src/main/java/io/milvus/v2/IteratorExample.java‎

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -160,6 +160,44 @@ private static void queryIterator(String expr, int batchSize, int offset, int li
160160
System.out.printf("%d query results returned%n", counter);
161161
}
162162

163+
private static void queryIteratorWithTemplate(int batchSize) {
164+
System.out.println("\n========== queryIterator() ==========");
165+
List<Long> ids = new ArrayList<>();
166+
for (long i = 500L; i < 600L; i++) {
167+
ids.add(i);
168+
}
169+
Map<String, Object> template = new HashMap<>();
170+
template.put("my_ids", ids);
171+
172+
String filter = ID_FIELD + " in {my_ids}";
173+
QueryIterator queryIterator = client.queryIterator(QueryIteratorReq.builder()
174+
.collectionName(COLLECTION_NAME)
175+
.expr(filter)
176+
.outputFields(Lists.newArrayList(ID_FIELD, AGE_FIELD))
177+
.batchSize(batchSize)
178+
.filterTemplateValues(template)
179+
.consistencyLevel(ConsistencyLevel.BOUNDED)
180+
.build());
181+
182+
System.out.println("QueryIterator with filter template results:");
183+
int counter = 0;
184+
while (true) {
185+
List<QueryResultsWrapper.RowRecord> res = queryIterator.next();
186+
if (res.isEmpty()) {
187+
System.out.println("query iteration finished, close");
188+
queryIterator.close();
189+
break;
190+
}
191+
192+
for (QueryResultsWrapper.RowRecord record : res) {
193+
System.out.println(record);
194+
counter++;
195+
}
196+
}
197+
System.out.printf("%d query results returned%n", counter);
198+
}
199+
200+
163201
// Search iterator V1
164202
private static void searchIteratorV1(String expr, String params, int batchSize, int topK) {
165203
System.out.println("\n========== searchIteratorV1() ==========");
@@ -235,16 +273,59 @@ private static void searchIteratorV2(String filter, Map<String, Object> params,
235273
System.out.printf("%d search results returned\n%n", counter);
236274
}
237275

276+
private static void searchIteratorV2WithTemplate(int batchSize) {
277+
System.out.println("\n========== searchIteratorV2() ==========");
278+
List<Long> ids = new ArrayList<>();
279+
for (long i = 500L; i < 600L; i++) {
280+
ids.add(i);
281+
}
282+
Map<String, Object> template = new HashMap<>();
283+
template.put("my_ids", ids);
284+
285+
String filter = ID_FIELD + " in {my_ids}";
286+
SearchIteratorV2 searchIterator = client.searchIteratorV2(SearchIteratorReqV2.builder()
287+
.collectionName(COLLECTION_NAME)
288+
.outputFields(Lists.newArrayList(AGE_FIELD))
289+
.batchSize(batchSize)
290+
.vectorFieldName(VECTOR_FIELD)
291+
.vectors(Collections.singletonList(new FloatVec(CommonUtils.generateFloatVector(VECTOR_DIM))))
292+
.filter(filter)
293+
.filterTemplateValues(template)
294+
.metricType(IndexParam.MetricType.L2)
295+
.consistencyLevel(ConsistencyLevel.BOUNDED)
296+
.build());
297+
298+
System.out.println("SearchIteratorV2 with filter template results:");
299+
int counter = 0;
300+
while (true) {
301+
List<SearchResp.SearchResult> res = searchIterator.next();
302+
if (res.isEmpty()) {
303+
System.out.println("Search iteration finished, close");
304+
searchIterator.close();
305+
break;
306+
}
307+
308+
for (SearchResp.SearchResult record : res) {
309+
System.out.println(record);
310+
counter++;
311+
}
312+
}
313+
System.out.printf("%d search results returned\n%n", counter);
314+
}
315+
238316
public static void main(String[] args) {
239317
buildCollection();
240318
queryIterator("userID < 300", 50, 5, 400);
319+
queryIteratorWithTemplate(80);
320+
241321
searchIteratorV1("userAge > 50 &&userAge < 100", "{\"range_filter\": 15.0, \"radius\": 20.0}", 100, 500);
242322
searchIteratorV1("", "", 10, 99);
243323
searchIteratorV2("userAge > 10 &&userAge < 20", null, 50, 120, null);
244324

245325
Map<String, Object> extraParams = new HashMap<>();
246326
extraParams.put("radius", 15.0);
247327
searchIteratorV2("", extraParams, 50, 100, null);
328+
searchIteratorV2WithTemplate(80);
248329

249330
// use external function to filter the result
250331
Function<List<SearchResp.SearchResult>, List<SearchResp.SearchResult>> externalFilterFunc = (List<SearchResp.SearchResult> src) -> {

‎sdk-core/src/main/java/io/milvus/orm/iterator/IteratorAdapterV2.java‎

Lines changed: 85 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,11 +27,13 @@
2727
import io.milvus.param.collection.FieldType;
2828
import io.milvus.param.dml.QueryIteratorParam;
2929
import io.milvus.param.dml.SearchIteratorParam;
30+
import io.milvus.v2.common.ConsistencyLevel;
3031
import io.milvus.v2.common.IndexParam;
3132
import io.milvus.v2.service.collection.request.CreateCollectionReq;
3233
import io.milvus.v2.service.vector.request.QueryIteratorReq;
3334
import io.milvus.v2.service.vector.request.SearchIteratorReq;
34-
import io.milvus.v2.service.vector.request.data.BaseVector;
35+
import io.milvus.v2.service.vector.request.data.*;
36+
import org.apache.commons.lang3.StringUtils;
3537

3638
import java.nio.ByteBuffer;
3739
import java.util.ArrayList;
@@ -59,6 +61,27 @@ public static QueryIteratorParam convertV2Req(QueryIteratorReq queryIteratorReq)
5961
return builder.build();
6062
}
6163

64+
public static QueryIteratorReq convertV1Param(QueryIteratorParam param) {
65+
ConsistencyLevel level = null;
66+
if (param.getConsistencyLevel() != null) {
67+
level = ConsistencyLevel.valueOf(param.getConsistencyLevel().name());
68+
}
69+
70+
return QueryIteratorReq.builder()
71+
.databaseName(param.getDatabaseName())
72+
.collectionName(param.getCollectionName())
73+
.partitionNames(param.getPartitionNames())
74+
.expr(param.getExpr())
75+
.outputFields(param.getOutFields())
76+
.offset(param.getOffset())
77+
.limit(param.getLimit())
78+
.ignoreGrowing(param.isIgnoreGrowing())
79+
.batchSize(param.getBatchSize())
80+
.reduceStopForBest(param.isReduceStopForBest())
81+
.consistencyLevel(level)
82+
.build();
83+
}
84+
6285
public static SearchIteratorParam convertV2Req(SearchIteratorReq searchIteratorReq) {
6386
MetricType metricType = MetricType.None;
6487
if (searchIteratorReq.getMetricType() != IndexParam.MetricType.INVALID) {
@@ -130,6 +153,67 @@ public static SearchIteratorParam convertV2Req(SearchIteratorReq searchIteratorR
130153
return builder.build();
131154
}
132155

156+
public static SearchIteratorReq convertV1Param(SearchIteratorParam param) {
157+
ConsistencyLevel level = null;
158+
if (param.getConsistencyLevel() != null) {
159+
level = ConsistencyLevel.valueOf(param.getConsistencyLevel().name());
160+
}
161+
162+
IndexParam.MetricType metricType = IndexParam.MetricType.INVALID;
163+
if (StringUtils.isEmpty(param.getMetricType())) {
164+
metricType = IndexParam.MetricType.valueOf(param.getMetricType());
165+
}
166+
167+
List<BaseVector> vectors = new ArrayList<>();
168+
switch (param.getPlType()) {
169+
case FloatVector: {
170+
List<List<Float>> data = (List<List<Float>>) param.getVectors();
171+
data.forEach(vector -> vectors.add(new FloatVec(vector)));
172+
break;
173+
}
174+
case BinaryVector: {
175+
List<ByteBuffer> data = (List<ByteBuffer>) param.getVectors();
176+
data.forEach(vector -> vectors.add(new BinaryVec(vector)));
177+
break;
178+
}
179+
case Float16Vector: {
180+
List<ByteBuffer> data = (List<ByteBuffer>) param.getVectors();
181+
data.forEach(vector -> vectors.add(new Float16Vec(vector)));
182+
break;
183+
}
184+
case BFloat16Vector: {
185+
List<ByteBuffer> data = (List<ByteBuffer>) param.getVectors();
186+
data.forEach(vector -> vectors.add(new BFloat16Vec(vector)));
187+
break;
188+
}
189+
case SparseFloatVector: {
190+
List<SortedMap<Long, Float>> data = (List<SortedMap<Long, Float>>) param.getVectors();
191+
data.forEach(vector -> vectors.add(new SparseFloatVec(vector)));
192+
break;
193+
}
194+
default:
195+
throw new ParamException("Unsupported vector type.");
196+
}
197+
198+
return SearchIteratorReq.builder()
199+
.databaseName(param.getDatabaseName())
200+
.collectionName(param.getCollectionName())
201+
.partitionNames(param.getPartitionNames())
202+
.vectorFieldName(param.getVectorFieldName())
203+
.vectors(vectors)
204+
.consistencyLevel(level)
205+
.metricType(metricType)
206+
.limit(param.getTopK())
207+
.expr(param.getExpr())
208+
.outputFields(param.getOutFields())
209+
.roundDecimal(param.getRoundDecimal())
210+
.params(param.getParams())
211+
.groupByFieldName(param.getGroupByFieldName())
212+
.ignoreGrowing(param.isIgnoreGrowing())
213+
.batchSize(param.getBatchSize())
214+
.build();
215+
}
216+
133217
public static FieldType convertV2Field(CreateCollectionReq.FieldSchema schema) {
134218
FieldType.Builder builder = FieldType.newBuilder()
135219
.withName(schema.getName())

‎sdk-core/src/main/java/io/milvus/orm/iterator/QueryIterator.java‎

Lines changed: 27 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -21,14 +21,14 @@
2121

2222
import io.milvus.grpc.*;
2323
import io.milvus.param.Constant;
24-
import io.milvus.param.ParamUtils;
2524
import io.milvus.param.collection.FieldType;
2625
import io.milvus.param.dml.QueryIteratorParam;
27-
import io.milvus.param.dml.QueryParam;
2826
import io.milvus.response.QueryResultsWrapper;
2927
import io.milvus.v2.service.collection.request.CreateCollectionReq;
3028
import io.milvus.v2.service.vector.request.QueryIteratorReq;
29+
import io.milvus.v2.service.vector.request.QueryReq;
3130
import io.milvus.v2.utils.RpcUtils;
31+
import io.milvus.v2.utils.VectorUtils;
3232
import org.apache.commons.lang3.StringUtils;
3333
import org.slf4j.Logger;
3434
import org.slf4j.LoggerFactory;
@@ -44,7 +44,7 @@ public class QueryIterator {
4444
private final MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub;
4545
private final FieldType primaryField;
4646

47-
private final QueryIteratorParam queryIteratorParam;
47+
private final QueryIteratorReq queryIteratorReq;
4848
private final int batchSize;
4949
private final long limit;
5050
private final String expr;
@@ -61,7 +61,7 @@ public QueryIterator(QueryIteratorParam queryIteratorParam,
6161
this.iteratorCache = new IteratorCache();
6262
this.blockingStub = blockingStub;
6363
this.primaryField = primaryField;
64-
this.queryIteratorParam = queryIteratorParam;
64+
this.queryIteratorReq = IteratorAdapterV2.convertV1Param(queryIteratorParam);
6565

6666
this.batchSize = (int) queryIteratorParam.getBatchSize();
6767
this.expr = queryIteratorParam.getExpr();
@@ -78,15 +78,14 @@ public QueryIterator(QueryIteratorReq queryIteratorReq,
7878
CreateCollectionReq.FieldSchema primaryField) {
7979
this.iteratorCache = new IteratorCache();
8080
this.blockingStub = blockingStub;
81-
IteratorAdapterV2 adapter = new IteratorAdapterV2();
82-
this.queryIteratorParam = adapter.convertV2Req(queryIteratorReq);
83-
this.primaryField = adapter.convertV2Field(primaryField);
81+
this.queryIteratorReq = queryIteratorReq;
82+
this.primaryField = IteratorAdapterV2.convertV2Field(primaryField);
8483

8584

86-
this.batchSize = (int) queryIteratorParam.getBatchSize();
87-
this.expr = queryIteratorParam.getExpr();
88-
this.limit = queryIteratorParam.getLimit();
89-
this.offset = queryIteratorParam.getOffset();
85+
this.batchSize = (int) queryIteratorReq.getBatchSize();
86+
this.expr = queryIteratorReq.getExpr();
87+
this.limit = queryIteratorReq.getLimit();
88+
this.offset = queryIteratorReq.getOffset();
9089
this.rpcUtils = new RpcUtils();
9190

9291
setupTsByRequest();
@@ -208,24 +207,27 @@ private boolean isResSufficient(List<QueryResultsWrapper.RowRecord> ret) {
208207
private QueryResults executeQuery(String expr, long offset, long limit, long ts, boolean isSeek) {
209208
// for seeking offset, no need to return output fields
210209
List<String> outputFields = new ArrayList<>();
211-
boolean reduceStopForBest = queryIteratorParam.isReduceStopForBest();
210+
boolean reduceStopForBest = queryIteratorReq.isReduceStopForBest();
212211
if (!isSeek) {
213-
outputFields = queryIteratorParam.getOutFields();
212+
outputFields = queryIteratorReq.getOutputFields();
214213
reduceStopForBest = false;
215214
}
216-
QueryParam queryParam = QueryParam.newBuilder()
217-
.withDatabaseName(queryIteratorParam.getDatabaseName())
218-
.withCollectionName(queryIteratorParam.getCollectionName())
219-
.withConsistencyLevel(queryIteratorParam.getConsistencyLevel())
220-
.withPartitionNames(queryIteratorParam.getPartitionNames())
221-
.withOutFields(outputFields)
222-
.withExpr(expr)
223-
.withOffset(offset)
224-
.withLimit(limit)
225-
.withIgnoreGrowing(queryIteratorParam.isIgnoreGrowing())
215+
QueryReq queryReq = QueryReq.builder()
216+
.databaseName(queryIteratorReq.getDatabaseName())
217+
.collectionName(queryIteratorReq.getCollectionName())
218+
.partitionNames(queryIteratorReq.getPartitionNames())
219+
.consistencyLevel(queryIteratorReq.getConsistencyLevel())
220+
.outputFields(outputFields)
221+
.filter(expr)
222+
.offset(offset)
223+
.limit(limit)
224+
.ignoreGrowing(queryIteratorReq.isIgnoreGrowing())
225+
.timezone(queryIteratorReq.getTimezone())
226+
.filterTemplateValues(queryIteratorReq.getFilterTemplateValues())
226227
.build();
227228

228-
QueryRequest queryRequest = ParamUtils.convertQueryParam(queryParam);
229+
VectorUtils vectorUtils = new VectorUtils();
230+
QueryRequest queryRequest = vectorUtils.ConvertToGrpcQueryRequest(queryReq);
229231
QueryRequest.Builder builder = queryRequest.toBuilder();
230232
// reduce stop for best
231233
builder.addQueryParams(KeyValuePair.newBuilder()
@@ -246,7 +248,7 @@ private QueryResults executeQuery(String expr, long offset, long limit, long ts,
246248
builder.setUseDefaultConsistency(true);
247249

248250
QueryResults response = rpcUtils.retry(() -> blockingStub.query(builder.build()));
249-
String title = String.format("QueryRequest collectionName:%s", queryIteratorParam.getCollectionName());
251+
String title = String.format("QueryRequest collectionName:%s", queryIteratorReq.getCollectionName());
250252
rpcUtils.handleResponse(title, response.getStatus());
251253
return response;
252254
}

‎sdk-core/src/main/java/io/milvus/orm/iterator/SearchIteratorV2.java‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@
4242
import static io.milvus.param.Constant.UNLIMITED;
4343

4444
public class SearchIteratorV2 {
45-
private static final Logger logger = LoggerFactory.getLogger(SearchIterator.class);
45+
private static final Logger logger = LoggerFactory.getLogger(SearchIteratorV2.class);
4646
private final MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub;
4747

4848
private final SearchIteratorReqV2 searchIteratorReq;
@@ -86,12 +86,12 @@ private void checkParams() {
8686
int rows = searchIteratorReq.getVectors().size();
8787
if (rows > 1) {
8888
ExceptionUtils.throwUnExpectedException("SearchIterator does not support processing multiple vectors simultaneously");
89-
} else if (rows <= 0) {
89+
} else if (rows == 0) {
9090
ExceptionUtils.throwUnExpectedException("The vector data for search cannot be empty");
9191
}
9292

93-
if (searchIteratorReq.getTopK() != UNLIMITED) {
94-
this.leftResCnt = searchIteratorReq.getTopK();
93+
if (searchIteratorReq.getLimit() != UNLIMITED) {
94+
this.leftResCnt = (int) searchIteratorReq.getLimit();
9595
}
9696
}
9797

@@ -117,15 +117,17 @@ private SearchResults executeSearch(int limit) {
117117
.databaseName(searchIteratorReq.getDatabaseName())
118118
.annsField(searchIteratorReq.getVectorFieldName())
119119
.data(searchIteratorReq.getVectors())
120-
.topK(limit)
120+
.limit(limit)
121121
.filter(searchIteratorReq.getFilter())
122122
.consistencyLevel(searchIteratorReq.getConsistencyLevel())
123123
.outputFields(searchIteratorReq.getOutputFields())
124124
.roundDecimal(searchIteratorReq.getRoundDecimal())
125125
.searchParams(searchParams)
126126
.metricType(searchIteratorReq.getMetricType())
127+
.timezone(searchIteratorReq.getTimezone())
127128
.ignoreGrowing(searchIteratorReq.isIgnoreGrowing())
128129
.groupByFieldName(searchIteratorReq.getGroupByFieldName())
130+
.filterTemplateValues(searchIteratorReq.getFilterTemplateValues())
129131
.build();
130132
SearchRequest searchRequest = new VectorUtils().ConvertToGrpcSearchRequest(request);
131133
SearchResults response = rpcUtils.retry(() -> this.blockingStub.search(searchRequest));

0 commit comments

Comments
 (0)