Skip to content

Commit bcf6bc1

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 bcf6bc1

5 files changed

Lines changed: 548 additions & 36 deletions

File tree

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

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

‎sdk-bulkwriter/src/main/java/io/milvus/bulkwriter/common/utils/ParquetUtils.java‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -201,6 +201,14 @@ private static void fillStructType(Types.MessageTypeBuilder messageTypeBuilder,
201201
case FloatVector:
202202
setMessageType(groupBuilder, PrimitiveType.PrimitiveTypeName.FLOAT, null, subField, true);
203203
break;
204+
case BinaryVector:
205+
case Float16Vector:
206+
case BFloat16Vector:
207+
case Int8Vector:
208+
boolean isSigned = (subField.getDataType() == io.milvus.v2.common.DataType.Int8Vector);
209+
setMessageType(groupBuilder, PrimitiveType.PrimitiveTypeName.INT32,
210+
LogicalTypeAnnotation.IntLogicalTypeAnnotation.intType(8, isSigned), subField, true);
211+
break;
204212
default:
205213
break;
206214
}

0 commit comments

Comments
 (0)