Skip to content

Commit 2982741

Browse files
committed
Support OrderBy for search/query
Signed-off-by: yhmo <yihua.mo@zilliz.com>
1 parent 4ad9742 commit 2982741

8 files changed

Lines changed: 480 additions & 0 deletions

File tree

Lines changed: 275 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,275 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package io.milvus.v2;
21+
22+
import com.google.gson.Gson;
23+
import com.google.gson.JsonArray;
24+
import com.google.gson.JsonObject;
25+
import io.milvus.v2.client.ConnectConfig;
26+
import io.milvus.v2.client.MilvusClientV2;
27+
import io.milvus.v2.common.ConsistencyLevel;
28+
import io.milvus.v2.common.DataType;
29+
import io.milvus.v2.common.IndexParam;
30+
import io.milvus.v2.service.collection.request.AddFieldReq;
31+
import io.milvus.v2.service.collection.request.CreateCollectionReq;
32+
import io.milvus.v2.service.collection.request.DropCollectionReq;
33+
import io.milvus.v2.service.vector.request.InsertReq;
34+
import io.milvus.v2.service.vector.request.QueryReq;
35+
import io.milvus.v2.service.vector.request.SearchReq;
36+
import io.milvus.v2.service.vector.request.aggregation.AggDirection;
37+
import io.milvus.v2.service.vector.request.aggregation.OrderByField;
38+
import io.milvus.v2.service.vector.request.data.FloatVec;
39+
import io.milvus.v2.service.vector.response.QueryResp;
40+
import io.milvus.v2.service.vector.response.SearchResp;
41+
42+
import java.util.ArrayList;
43+
import java.util.Arrays;
44+
import java.util.Collections;
45+
import java.util.List;
46+
import java.util.Random;
47+
48+
public class OrderByExample {
49+
private static final String COLLECTION_NAME = "java_sdk_example_order_by_v2";
50+
private static final String ID_FIELD = "id";
51+
private static final String PRICE_FIELD = "price";
52+
private static final String RATING_FIELD = "rating";
53+
private static final String CATEGORY_FIELD = "category";
54+
private static final String METADATA_FIELD = "metadata";
55+
private static final String VECTOR_FIELD = "embeddings";
56+
private static final String DYNAMIC_VIEWS_FIELD = "dynamic_views";
57+
private static final int DIM = 8;
58+
private static final int NUM_ENTITIES = 200;
59+
private static final Random RNG = new Random(19530L);
60+
61+
private static List<OrderByField> orderByFields(OrderByField... fields) {
62+
return Arrays.asList(fields);
63+
}
64+
65+
private static void printSearchResults(SearchResp searchResp, String title) {
66+
System.out.printf("%n============================================================%n%s%n============================================================%n", title);
67+
List<List<SearchResp.SearchResult>> searchResults = searchResp.getSearchResults();
68+
for (int i = 0; i < searchResults.size(); i++) {
69+
System.out.printf("Search result %d:%n", i);
70+
for (SearchResp.SearchResult result : searchResults.get(i)) {
71+
System.out.println(result);
72+
}
73+
}
74+
}
75+
76+
private static void printQueryResults(QueryResp queryResp, String title) {
77+
System.out.printf("%n============================================================%n%s%n============================================================%n", title);
78+
for (QueryResp.QueryResult result : queryResp.getQueryResults()) {
79+
System.out.println(result.getEntity());
80+
}
81+
}
82+
83+
private static float[] randomVector() {
84+
float[] values = new float[DIM];
85+
for (int i = 0; i < DIM; i++) {
86+
values[i] = RNG.nextFloat();
87+
}
88+
return values;
89+
}
90+
91+
private static void prepareCollection(MilvusClientV2 client) {
92+
client.dropCollection(DropCollectionReq.builder()
93+
.collectionName(COLLECTION_NAME)
94+
.build());
95+
96+
CreateCollectionReq.CollectionSchema schema = CreateCollectionReq.CollectionSchema.builder()
97+
.enableDynamicField(true)
98+
.build();
99+
schema.addField(AddFieldReq.builder()
100+
.fieldName(ID_FIELD)
101+
.dataType(DataType.Int64)
102+
.isPrimaryKey(true)
103+
.autoID(false)
104+
.build());
105+
schema.addField(AddFieldReq.builder()
106+
.fieldName(PRICE_FIELD)
107+
.dataType(DataType.Double)
108+
.build());
109+
schema.addField(AddFieldReq.builder()
110+
.fieldName(RATING_FIELD)
111+
.dataType(DataType.Double)
112+
.build());
113+
schema.addField(AddFieldReq.builder()
114+
.fieldName(CATEGORY_FIELD)
115+
.dataType(DataType.VarChar)
116+
.maxLength(64)
117+
.build());
118+
schema.addField(AddFieldReq.builder()
119+
.fieldName(METADATA_FIELD)
120+
.dataType(DataType.JSON)
121+
.build());
122+
schema.addField(AddFieldReq.builder()
123+
.fieldName(VECTOR_FIELD)
124+
.dataType(DataType.FloatVector)
125+
.dimension(DIM)
126+
.build());
127+
128+
List<IndexParam> indexes = new ArrayList<>();
129+
indexes.add(IndexParam.builder()
130+
.fieldName(VECTOR_FIELD)
131+
.indexType(IndexParam.IndexType.IVF_FLAT)
132+
.metricType(IndexParam.MetricType.L2)
133+
.extraParams(Collections.singletonMap("nlist", 128))
134+
.build());
135+
136+
client.createCollection(CreateCollectionReq.builder()
137+
.collectionName(COLLECTION_NAME)
138+
.collectionSchema(schema)
139+
.indexParams(indexes)
140+
.consistencyLevel(ConsistencyLevel.BOUNDED)
141+
.build());
142+
System.out.println("Collection created");
143+
}
144+
145+
private static void insertRows(MilvusClientV2 client) {
146+
List<JsonObject> rows = new ArrayList<>();
147+
Gson gson = new Gson();
148+
List<String> categories = Arrays.asList("cate1", "cate2", "cate3", "cate4", "cate5");
149+
for (int i = 0; i < NUM_ENTITIES; i++) {
150+
JsonObject row = new JsonObject();
151+
row.addProperty(ID_FIELD, i);
152+
row.addProperty(PRICE_FIELD, 10.0 + (i % 12));
153+
row.addProperty(RATING_FIELD, Math.round(RNG.nextDouble() * 50.0) / 10.0);
154+
row.addProperty(CATEGORY_FIELD, categories.get(i % categories.size()));
155+
row.add(VECTOR_FIELD, gson.toJsonTree(randomVector()));
156+
157+
JsonObject metadata = new JsonObject();
158+
metadata.addProperty("age", 18 + (i % 40));
159+
metadata.addProperty("score", i % 101);
160+
metadata.addProperty("popularity", Math.round(RNG.nextDouble() * 100.0) / 10.0);
161+
JsonArray tags = new JsonArray();
162+
tags.add(categories.get(i % categories.size()));
163+
tags.add("tag_" + (i % 10));
164+
metadata.add("tags", tags);
165+
row.add(METADATA_FIELD, metadata);
166+
167+
row.addProperty(DYNAMIC_VIEWS_FIELD, i * 10 + RNG.nextInt(100));
168+
rows.add(row);
169+
}
170+
171+
client.insert(InsertReq.builder()
172+
.collectionName(COLLECTION_NAME)
173+
.data(rows)
174+
.build());
175+
System.out.printf("%d rows inserted%n", rows.size());
176+
}
177+
178+
private static void waitForVisible(MilvusClientV2 client) {
179+
QueryResp countResp = client.query(QueryReq.builder()
180+
.collectionName(COLLECTION_NAME)
181+
.outputFields(Collections.singletonList("count(*)"))
182+
.consistencyLevel(ConsistencyLevel.STRONG)
183+
.build());
184+
System.out.printf("%d rows persisted%n", (long) countResp.getQueryResults().get(0).getEntity().get("count(*)"));
185+
}
186+
187+
private static void searchExamples(MilvusClientV2 client) {
188+
List<FloatVec> vectors = Collections.singletonList(new FloatVec(randomVector()));
189+
190+
SearchResp searchByPrice = client.search(SearchReq.builder()
191+
.collectionName(COLLECTION_NAME)
192+
.annsField(VECTOR_FIELD)
193+
.data(new ArrayList<>(vectors))
194+
.limit(10)
195+
.outputFields(Arrays.asList(ID_FIELD, PRICE_FIELD, RATING_FIELD, CATEGORY_FIELD))
196+
.orderByFields(orderByFields(
197+
OrderByField.builder().fieldName(PRICE_FIELD).direction(AggDirection.ASC).build()))
198+
.build());
199+
printSearchResults(searchByPrice, "Search with order_by_fields price ASC");
200+
201+
SearchResp searchByPriceAndRating = client.search(SearchReq.builder()
202+
.collectionName(COLLECTION_NAME)
203+
.annsField(VECTOR_FIELD)
204+
.data(new ArrayList<>(vectors))
205+
.limit(10)
206+
.outputFields(Arrays.asList(ID_FIELD, PRICE_FIELD, RATING_FIELD, CATEGORY_FIELD, METADATA_FIELD))
207+
.orderByFields(orderByFields(
208+
OrderByField.builder().fieldName(PRICE_FIELD).direction(AggDirection.ASC).build(),
209+
OrderByField.builder().fieldName(RATING_FIELD).direction(AggDirection.DESC).build()))
210+
.build());
211+
printSearchResults(searchByPriceAndRating, "Search with order_by_fields price ASC, rating DESC");
212+
213+
SearchResp searchByJsonAndDynamic = client.search(SearchReq.builder()
214+
.collectionName(COLLECTION_NAME)
215+
.annsField(VECTOR_FIELD)
216+
.data(new ArrayList<>(vectors))
217+
.limit(10)
218+
.outputFields(Arrays.asList(ID_FIELD, PRICE_FIELD, RATING_FIELD, CATEGORY_FIELD, METADATA_FIELD, DYNAMIC_VIEWS_FIELD))
219+
.orderByFields(orderByFields(
220+
OrderByField.builder().fieldName("metadata[\"age\"]").direction(AggDirection.ASC).build(),
221+
OrderByField.builder().fieldName(DYNAMIC_VIEWS_FIELD).direction(AggDirection.DESC).build()))
222+
.build());
223+
printSearchResults(searchByJsonAndDynamic, "Search with order_by_fields metadata[\"age\"] ASC, dynamic_views DESC");
224+
}
225+
226+
private static void queryExamples(MilvusClientV2 client) {
227+
QueryResp queryByPrice = client.query(QueryReq.builder()
228+
.collectionName(COLLECTION_NAME)
229+
.filter(ID_FIELD + " >= 0")
230+
.limit(30)
231+
.outputFields(Arrays.asList(ID_FIELD, PRICE_FIELD, RATING_FIELD, CATEGORY_FIELD))
232+
.orderByFields(orderByFields(
233+
OrderByField.builder().fieldName(PRICE_FIELD).direction(AggDirection.DESC).build()))
234+
.build());
235+
printQueryResults(queryByPrice, "Query with order_by_fields price DESC");
236+
237+
QueryResp queryByPriceAndRating = client.query(QueryReq.builder()
238+
.collectionName(COLLECTION_NAME)
239+
.filter(ID_FIELD + " >= 0")
240+
.limit(30)
241+
.outputFields(Arrays.asList(ID_FIELD, PRICE_FIELD, RATING_FIELD, CATEGORY_FIELD, METADATA_FIELD))
242+
.orderByFields(orderByFields(
243+
OrderByField.builder().fieldName(PRICE_FIELD).direction(AggDirection.ASC).build(),
244+
OrderByField.builder().fieldName(RATING_FIELD).direction(AggDirection.DESC).build()))
245+
.build());
246+
printQueryResults(queryByPriceAndRating, "Query with order_by_fields price ASC, rating DESC");
247+
248+
QueryResp queryByCategoryAndPrice = client.query(QueryReq.builder()
249+
.collectionName(COLLECTION_NAME)
250+
.filter(ID_FIELD + " >= 0")
251+
.limit(30)
252+
.outputFields(Arrays.asList(ID_FIELD, PRICE_FIELD, RATING_FIELD, CATEGORY_FIELD, METADATA_FIELD, DYNAMIC_VIEWS_FIELD))
253+
.orderByFields(orderByFields(
254+
OrderByField.builder().fieldName(CATEGORY_FIELD).direction(AggDirection.ASC).build(),
255+
OrderByField.builder().fieldName(PRICE_FIELD).direction(AggDirection.DESC).build()))
256+
.build());
257+
printQueryResults(queryByCategoryAndPrice, "Query with order_by_fields category ASC, price DESC");
258+
}
259+
260+
public static void main(String[] args) {
261+
ConnectConfig config = ConnectConfig.builder()
262+
.uri("http://localhost:19530")
263+
.build();
264+
MilvusClientV2 client = new MilvusClientV2(config);
265+
266+
prepareCollection(client);
267+
insertRows(client);
268+
waitForVisible(client);
269+
searchExamples(client);
270+
queryExamples(client);
271+
272+
client.close();
273+
System.out.println("All order_by examples completed!");
274+
}
275+
}

sdk-core/src/main/java/io/milvus/param/Constant.java

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ public class Constant {
3636
public static final String ITERATOR_FIELD = "iterator";
3737
public static final String CLUSTER_ID = "cluster_id";
3838
public static final String GROUP_BY_FIELD = "group_by_field";
39+
public static final String ORDER_BY_FIELDS = "order_by_fields";
3940
public static final String GROUP_SIZE = "group_size";
4041
public static final String STRICT_GROUP_SIZE = "strict_group_size";
4142
public static final String HINTS = "hints"; // from milvus v2.5

sdk-core/src/main/java/io/milvus/v2/client/MilvusClientV2Session.java

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,7 @@ private SearchReq copy(SearchReq request) {
113113
.consistencyLevel(request.getConsistencyLevel())
114114
.ignoreGrowing(request.isIgnoreGrowing())
115115
.timezone(request.getTimezone())
116+
.orderByFields(request.getOrderByFields())
116117
.groupByFieldName(request.getGroupByFieldName())
117118
.groupSize(request.getGroupSize())
118119
.strictGroupSize(request.getStrictGroupSize())
@@ -160,6 +161,7 @@ private QueryReq copy(QueryReq request) {
160161
.limit(request.getLimit())
161162
.ignoreGrowing(request.isIgnoreGrowing())
162163
.timezone(request.getTimezone())
164+
.orderByFields(request.getOrderByFields())
163165
.queryParams(request.getQueryParams())
164166
.filterTemplateValues(request.getFilterTemplateValues())
165167
.build();

sdk-core/src/main/java/io/milvus/v2/service/vector/request/QueryReq.java

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
package io.milvus.v2.service.vector.request;
2121

2222
import io.milvus.v2.common.ConsistencyLevel;
23+
import io.milvus.v2.service.vector.request.aggregation.OrderByField;
2324

2425
import java.util.*;
2526

@@ -36,6 +37,7 @@ public class QueryReq {
3637
private long limit;
3738
private boolean ignoreGrowing;
3839
private String timezone;
40+
private List<OrderByField> orderByFields;
3941

4042
// Extra parameters for query, timezone, time_fields, etc.
4143
// Make sure the value can be converted to String by String.valueOf().
@@ -64,6 +66,7 @@ private QueryReq(QueryReqBuilder builder) {
6466
this.offset = builder.offset;
6567
this.limit = builder.limit;
6668
this.ignoreGrowing = builder.ignoreGrowing;
69+
this.orderByFields = builder.orderByFields;
6770
this.queryParams = builder.queryParams;
6871
this.filterTemplateValues = builder.filterTemplateValues;
6972
this.timezone = builder.timezone;
@@ -165,6 +168,14 @@ public String getTimezone() {
165168
return timezone;
166169
}
167170

171+
public List<OrderByField> getOrderByFields() {
172+
return orderByFields;
173+
}
174+
175+
public void setOrderByFields(List<OrderByField> orderByFields) {
176+
this.orderByFields = orderByFields;
177+
}
178+
168179
public Map<String, Object> getQueryParams() {
169180
return queryParams;
170181
}
@@ -196,6 +207,7 @@ public String toString() {
196207
", limit=" + limit +
197208
", ignoreGrowing=" + ignoreGrowing +
198209
", timezone='" + timezone + '\'' +
210+
", orderByFields=" + orderByFields +
199211
", queryParams=" + queryParams +
200212
// ", filterTemplateValues=" + filterTemplateValues +
201213
'}';
@@ -214,6 +226,7 @@ public static class QueryReqBuilder {
214226
private long limit;
215227
private boolean ignoreGrowing;
216228
private String timezone = "";
229+
private List<OrderByField> orderByFields = new ArrayList<>();
217230
private Map<String, Object> queryParams = new HashMap<>();
218231
private Map<String, Object> filterTemplateValues = new HashMap<>();
219232

@@ -277,6 +290,11 @@ public QueryReqBuilder timezone(String timezone) {
277290
return this;
278291
}
279292

293+
public QueryReqBuilder orderByFields(List<OrderByField> orderByFields) {
294+
this.orderByFields = orderByFields;
295+
return this;
296+
}
297+
280298
public QueryReqBuilder queryParams(Map<String, Object> queryParams) {
281299
this.queryParams = queryParams;
282300
return this;

0 commit comments

Comments
 (0)