Skip to content

Commit 2b52a0f

Browse files
yhmoyhmo
committed
BulkWriter support binary/float16/int8 vectors in Struct field
Signed-off-by: yhmo <yihua.mo@zilliz.com>
1 parent 46b3538 commit 2b52a0f

5 files changed

Lines changed: 593 additions & 36 deletions

File tree

‎examples/src/main/java/io/milvus/v2/bulkwriter/BulkWriterStructExample.java‎

Lines changed: 389 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,389 @@
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+
package io.milvus.v2.bulkwriter;
20+
21+
import com.google.gson.Gson;
22+
import com.google.gson.JsonArray;
23+
import com.google.gson.JsonObject;
24+
import io.milvus.bulkwriter.RemoteBulkWriter;
25+
import io.milvus.bulkwriter.RemoteBulkWriterParam;
26+
import io.milvus.bulkwriter.common.clientenum.BulkFileType;
27+
import io.milvus.bulkwriter.common.clientenum.CloudStorage;
28+
import io.milvus.bulkwriter.connect.AzureConnectParam;
29+
import io.milvus.bulkwriter.connect.S3ConnectParam;
30+
import io.milvus.bulkwriter.connect.StorageConnectParam;
31+
import io.milvus.bulkwriter.request.describe.MilvusDescribeImportRequest;
32+
import io.milvus.bulkwriter.request.import_.MilvusImportRequest;
33+
import io.milvus.bulkwriter.restful.BulkImportUtils;
34+
import io.milvus.v1.CommonUtils;
35+
import io.milvus.v2.client.ConnectConfig;
36+
import io.milvus.v2.client.MilvusClientV2;
37+
import io.milvus.v2.common.ConsistencyLevel;
38+
import io.milvus.v2.common.DataType;
39+
import io.milvus.v2.common.IndexParam;
40+
import io.milvus.v2.service.collection.request.AddFieldReq;
41+
import io.milvus.v2.service.collection.request.CreateCollectionReq;
42+
import io.milvus.v2.service.collection.request.DropCollectionReq;
43+
import io.milvus.v2.service.collection.request.LoadCollectionReq;
44+
import io.milvus.v2.service.collection.request.RefreshLoadReq;
45+
import io.milvus.v2.service.database.request.CreateDatabaseReq;
46+
import io.milvus.v2.service.database.response.ListDatabasesResp;
47+
import io.milvus.v2.service.index.request.CreateIndexReq;
48+
import io.milvus.v2.service.vector.request.QueryReq;
49+
import io.milvus.v2.service.vector.response.QueryResp;
50+
51+
import java.io.IOException;
52+
import java.util.ArrayList;
53+
import java.util.Arrays;
54+
import java.util.List;
55+
import java.util.Map;
56+
import java.util.concurrent.TimeUnit;
57+
58+
public class BulkWriterStructExample {
59+
public static final String HOST = "127.0.0.1";
60+
public static final Integer PORT = 19530;
61+
public static final String USER_NAME = "user.name";
62+
public static final String PASSWORD = "password";
63+
64+
public static class StorageConsts {
65+
public static final CloudStorage cloudStorage = CloudStorage.MINIO;
66+
public static final String STORAGE_ENDPOINT = cloudStorage.getEndpoint("http://127.0.0.1:9000");
67+
public static final String STORAGE_BUCKET = "a-bucket";
68+
public static final String STORAGE_ACCESS_KEY = "minioadmin";
69+
public static final String STORAGE_SECRET_KEY = "minioadmin";
70+
public static final String STORAGE_REGION = "";
71+
public static final String AZURE_CONTAINER_NAME = "azure.container.name";
72+
public static final String AZURE_ACCOUNT_NAME = "azure.account.name";
73+
public static final String AZURE_ACCOUNT_KEY = "azure.account.key";
74+
}
75+
76+
private static final Gson GSON = new Gson();
77+
private static final String DATABASE_NAME = "java_sdk_db";
78+
private static final String COLLECTION_NAME = "java_sdk_bulkwriter_struct_v2";
79+
private static final String STRUCT_FIELD = "struct_field";
80+
private static final int DIM = 16;
81+
private static final int IVF_NLIST = 64;
82+
private static final BulkFileType FILE_TYPE = BulkFileType.PARQUET;
83+
private static MilvusClientV2 milvusClient;
84+
85+
public static void main(String[] args) throws Exception {
86+
createConnection();
87+
CreateCollectionReq.CollectionSchema schema = buildSchema();
88+
List<JsonObject> rows = buildRows(5);
89+
createCollection(schema);
90+
List<List<String>> batchFiles = writeRows(schema, rows);
91+
callBulkInsert(batchFiles);
92+
verifyImport(rows);
93+
}
94+
95+
private static void createConnection() {
96+
String url = String.format("http://%s:%s", HOST, PORT);
97+
milvusClient = new MilvusClientV2(ConnectConfig.builder()
98+
.uri(url)
99+
.username(USER_NAME)
100+
.password(PASSWORD)
101+
.build());
102+
103+
ListDatabasesResp dbs = milvusClient.listDatabases();
104+
if (!dbs.getDatabaseNames().contains(DATABASE_NAME)) {
105+
milvusClient.createDatabase(CreateDatabaseReq.builder()
106+
.databaseName(DATABASE_NAME)
107+
.build());
108+
}
109+
try {
110+
milvusClient.useDatabase(DATABASE_NAME);
111+
} catch (Exception e) {
112+
System.out.println("Unable to switch database, error: " + e);
113+
}
114+
}
115+
116+
private static CreateCollectionReq.CollectionSchema buildSchema() {
117+
CreateCollectionReq.CollectionSchema schema = CreateCollectionReq.CollectionSchema.builder()
118+
.build();
119+
schema.addField(AddFieldReq.builder()
120+
.fieldName("id")
121+
.dataType(DataType.Int64)
122+
.isPrimaryKey(true)
123+
.autoID(false)
124+
.build());
125+
schema.addField(AddFieldReq.builder()
126+
.fieldName("struct_field")
127+
.dataType(DataType.Array)
128+
.elementType(DataType.Struct)
129+
.maxCapacity(8)
130+
.addStructField(AddFieldReq.builder()
131+
.fieldName("st_name")
132+
.dataType(DataType.VarChar)
133+
.maxLength(128)
134+
.build())
135+
.addStructField(AddFieldReq.builder()
136+
.fieldName("st_float_vector")
137+
.dataType(DataType.FloatVector)
138+
.dimension(DIM)
139+
.build())
140+
.addStructField(AddFieldReq.builder()
141+
.fieldName("st_binary_vector")
142+
.dataType(DataType.BinaryVector)
143+
.dimension(DIM)
144+
.build())
145+
.addStructField(AddFieldReq.builder()
146+
.fieldName("st_float16_vector")
147+
.dataType(DataType.Float16Vector)
148+
.dimension(DIM)
149+
.build())
150+
.addStructField(AddFieldReq.builder()
151+
.fieldName("st_int8_vector")
152+
.dataType(DataType.Int8Vector)
153+
.dimension(DIM)
154+
.build())
155+
.build());
156+
return schema;
157+
}
158+
159+
private static List<JsonObject> buildRows(int rowCount) {
160+
List<JsonObject> rows = new ArrayList<>();
161+
for (int i = 0; i < rowCount; i++) {
162+
JsonObject row = new JsonObject();
163+
row.addProperty("id", i);
164+
165+
JsonArray structArray = new JsonArray();
166+
for (int j = 0; j < i % 3 + 1; j++) {
167+
structArray.add(buildStruct(i, j));
168+
}
169+
row.add("struct_field", structArray);
170+
rows.add(row);
171+
}
172+
return rows;
173+
}
174+
175+
private static JsonObject buildStruct(int rowId, int structId) {
176+
JsonObject struct = new JsonObject();
177+
struct.addProperty("st_name", String.format("row_%d_struct_%d", rowId, structId));
178+
179+
List<Float> floatVector = CommonUtils.generateFloatVector(DIM);
180+
List<Float> float16Vector = CommonUtils.generateFloatVector(DIM);
181+
struct.add("st_float_vector", GSON.toJsonTree(floatVector));
182+
struct.add("st_binary_vector", GSON.toJsonTree(CommonUtils.generateBinaryVector(DIM).array()));
183+
struct.add("st_float16_vector", GSON.toJsonTree(CommonUtils.encodeFloat16Vector(float16Vector, false).array()));
184+
struct.add("st_int8_vector", GSON.toJsonTree(CommonUtils.generateInt8Vector(DIM).array()));
185+
return struct;
186+
}
187+
188+
private static List<List<String>> writeRows(CreateCollectionReq.CollectionSchema schema, List<JsonObject> rows) throws Exception {
189+
try (RemoteBulkWriter remoteBulkWriter = buildRemoteBulkWriter(schema)) {
190+
for (JsonObject row : rows) {
191+
remoteBulkWriter.appendRow(row);
192+
}
193+
System.out.printf("%s rows appended%n", remoteBulkWriter.getTotalRowCount());
194+
remoteBulkWriter.commit(false);
195+
System.out.printf("Uploaded files: %s%n", remoteBulkWriter.getBatchFiles());
196+
return remoteBulkWriter.getBatchFiles();
197+
}
198+
}
199+
200+
private static RemoteBulkWriter buildRemoteBulkWriter(CreateCollectionReq.CollectionSchema schema) throws IOException {
201+
RemoteBulkWriterParam bulkWriterParam = RemoteBulkWriterParam.newBuilder()
202+
.withCollectionSchema(schema)
203+
.withRemotePath("bulk_data/struct_vector_example")
204+
.withFileType(FILE_TYPE)
205+
.withChunkSize(128 * 1024 * 1024)
206+
.withConnectParam(buildStorageConnectParam())
207+
.build();
208+
return new RemoteBulkWriter(bulkWriterParam);
209+
}
210+
211+
private static StorageConnectParam buildStorageConnectParam() {
212+
if (CloudStorage.isAzCloud(StorageConsts.cloudStorage.getCloudName())) {
213+
String connectionStr = "DefaultEndpointsProtocol=https;AccountName=" + StorageConsts.AZURE_ACCOUNT_NAME
214+
+ ";AccountKey=" + StorageConsts.AZURE_ACCOUNT_KEY
215+
+ ";EndpointSuffix=core.windows.net";
216+
return AzureConnectParam.newBuilder()
217+
.withConnStr(connectionStr)
218+
.withContainerName(StorageConsts.AZURE_CONTAINER_NAME)
219+
.build();
220+
}
221+
222+
return S3ConnectParam.newBuilder()
223+
.withEndpoint(StorageConsts.STORAGE_ENDPOINT)
224+
.withCloudName(StorageConsts.cloudStorage.getCloudName())
225+
.withBucketName(StorageConsts.STORAGE_BUCKET)
226+
.withAccessKey(StorageConsts.STORAGE_ACCESS_KEY)
227+
.withSecretKey(StorageConsts.STORAGE_SECRET_KEY)
228+
.withRegion(StorageConsts.STORAGE_REGION)
229+
.build();
230+
}
231+
232+
private static void createCollection(CreateCollectionReq.CollectionSchema schema) {
233+
milvusClient.dropCollection(DropCollectionReq.builder()
234+
.collectionName(COLLECTION_NAME)
235+
.databaseName(DATABASE_NAME)
236+
.build());
237+
milvusClient.createCollection(CreateCollectionReq.builder()
238+
.collectionName(COLLECTION_NAME)
239+
.databaseName(DATABASE_NAME)
240+
.collectionSchema(schema)
241+
.consistencyLevel(ConsistencyLevel.BOUNDED)
242+
.build());
243+
System.out.printf("Collection %s created%n", COLLECTION_NAME);
244+
}
245+
246+
private static void callBulkInsert(List<List<String>> batchFiles) throws InterruptedException {
247+
String url = String.format("http://%s:%s", HOST, PORT);
248+
MilvusImportRequest request = MilvusImportRequest.builder()
249+
.collectionName(COLLECTION_NAME)
250+
.dbName(DATABASE_NAME)
251+
.files(batchFiles)
252+
.apiKey(USER_NAME + ":" + PASSWORD)
253+
.build();
254+
String importResult = BulkImportUtils.bulkImport(url, request);
255+
System.out.println(importResult);
256+
257+
String jobId = GSON.fromJson(importResult, JsonObject.class)
258+
.getAsJsonObject("data")
259+
.get("jobId")
260+
.getAsString();
261+
262+
while (true) {
263+
TimeUnit.SECONDS.sleep(5);
264+
String progressResult = BulkImportUtils.getImportProgress(url, MilvusDescribeImportRequest.builder()
265+
.jobId(jobId)
266+
.apiKey(USER_NAME + ":" + PASSWORD)
267+
.build());
268+
JsonObject progress = GSON.fromJson(progressResult, JsonObject.class);
269+
String state = progress.getAsJsonObject("data").get("state").getAsString();
270+
if ("Failed".equals(state)) {
271+
String reason = progress.getAsJsonObject("data").get("reason").getAsString();
272+
throw new RuntimeException(String.format("Import job %s failed: %s", jobId, reason));
273+
}
274+
if ("Completed".equals(state)) {
275+
System.out.printf("Import job %s completed%n", jobId);
276+
break;
277+
}
278+
System.out.printf("Import job %s is running, state=%s%n", jobId, state);
279+
}
280+
}
281+
282+
private static void createIndexes() {
283+
List<IndexParam> indexParams = new ArrayList<>();
284+
indexParams.add(IndexParam.builder()
285+
.fieldName(String.format("%s[%s]", STRUCT_FIELD, "st_float_vector"))
286+
.indexName("index_st_float_vector")
287+
.indexType(IndexParam.IndexType.HNSW)
288+
.metricType(IndexParam.MetricType.MAX_SIM_COSINE)
289+
.build());
290+
indexParams.add(IndexParam.builder()
291+
.fieldName(String.format("%s[%s]", STRUCT_FIELD, "st_binary_vector"))
292+
.indexName("index_st_binary_vector")
293+
.indexType(IndexParam.IndexType.HNSW)
294+
.metricType(IndexParam.MetricType.MAX_SIM_HAMMING)
295+
.build());
296+
indexParams.add(IndexParam.builder()
297+
.fieldName(String.format("%s[%s]", STRUCT_FIELD, "st_float16_vector"))
298+
.indexName("index_st_float16_vector")
299+
.indexType(IndexParam.IndexType.IVF_FLAT)
300+
.metricType(IndexParam.MetricType.MAX_SIM_COSINE)
301+
.extraParams(buildIvfParams())
302+
.build());
303+
indexParams.add(IndexParam.builder()
304+
.fieldName(String.format("%s[%s]", STRUCT_FIELD, "st_int8_vector"))
305+
.indexName("index_st_int8_vector")
306+
.indexType(IndexParam.IndexType.HNSW)
307+
.metricType(IndexParam.MetricType.MAX_SIM_L2)
308+
.build());
309+
milvusClient.createIndex(CreateIndexReq.builder()
310+
.collectionName(COLLECTION_NAME)
311+
.databaseName(DATABASE_NAME)
312+
.indexParams(indexParams)
313+
.build());
314+
}
315+
316+
private static Map<String, Object> buildIvfParams() {
317+
Map<String, Object> params = new java.util.HashMap<>();
318+
params.put("nlist", IVF_NLIST);
319+
return params;
320+
}
321+
322+
private static void verifyImport(List<JsonObject> rows) {
323+
createIndexes();
324+
milvusClient.loadCollection(LoadCollectionReq.builder()
325+
.collectionName(COLLECTION_NAME)
326+
.databaseName(DATABASE_NAME)
327+
.build());
328+
milvusClient.refreshLoad(RefreshLoadReq.builder()
329+
.collectionName(COLLECTION_NAME)
330+
.databaseName(DATABASE_NAME)
331+
.build());
332+
333+
List<Long> ids = Arrays.asList(0L, (long) rows.size() - 1);
334+
QueryResp queryResp = milvusClient.query(QueryReq.builder()
335+
.collectionName(COLLECTION_NAME)
336+
.databaseName(DATABASE_NAME)
337+
.filter(String.format("id in %s", ids))
338+
.outputFields(Arrays.asList("*"))
339+
.consistencyLevel(ConsistencyLevel.STRONG)
340+
.build());
341+
342+
List<QueryResp.QueryResult> results = queryResp.getQueryResults();
343+
if (results.size() != ids.size()) {
344+
throw new RuntimeException("Unexpected query result count");
345+
}
346+
347+
for (QueryResp.QueryResult result : results) {
348+
Map<String, Object> entity = result.getEntity();
349+
long id = (Long) entity.get("id");
350+
JsonArray expectedStructs = rows.get((int) id).get("struct_field").getAsJsonArray();
351+
Object structField = entity.get("struct_field");
352+
if (!(structField instanceof List<?>)) {
353+
throw new RuntimeException("struct_field should be returned as a list");
354+
}
355+
356+
List<?> fetchedStructs = (List<?>) structField;
357+
if (fetchedStructs.size() != expectedStructs.size()) {
358+
throw new RuntimeException("Struct element count is unmatched");
359+
}
360+
361+
for (int i = 0; i < fetchedStructs.size(); i++) {
362+
if (!(fetchedStructs.get(i) instanceof Map<?, ?>)) {
363+
throw new RuntimeException("Struct element should be returned as a map");
364+
}
365+
366+
Map<?, ?> fetchedStruct = (Map<?, ?>) fetchedStructs.get(i);
367+
JsonObject expectedStruct = expectedStructs.get(i).getAsJsonObject();
368+
String expectedName = expectedStruct.get("st_name").getAsString();
369+
if (!expectedName.equals(fetchedStruct.get("st_name"))) {
370+
throw new RuntimeException("Struct field st_name is unmatched");
371+
}
372+
assertVectorFieldExists(fetchedStruct, "st_float_vector");
373+
assertVectorFieldExists(fetchedStruct, "st_binary_vector");
374+
assertVectorFieldExists(fetchedStruct, "st_float16_vector");
375+
assertVectorFieldExists(fetchedStruct, "st_int8_vector");
376+
}
377+
378+
System.out.printf("Verified row %d: %s%n", id, entity);
379+
}
380+
System.out.println("RemoteBulkWriter struct vector import is correct!");
381+
}
382+
383+
private static void assertVectorFieldExists(Map<?, ?> struct, String fieldName) {
384+
Object value = struct.get(fieldName);
385+
if (value == null) {
386+
throw new RuntimeException(String.format("Struct vector field '%s' is missing", fieldName));
387+
}
388+
}
389+
}

0 commit comments

Comments
 (0)