BulkWriter support binary/float16/int8 vectors in Struct field

Signed-off-by: yhmo <yihua.mo@zilliz.com>

Author: yhmo

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

@@ -190,7 +190,13 @@ protected void callBack(List<String> fileList) {
     @Override
     public void close() throws Exception {
         logger.info("execute remaining actions to prevent loss of memory data or residual empty directories.");
-        exit();
+        try {
+            exit();
+        } finally {
+            if (storageClient != null) {
+                storageClient.close();
+            }
+        }
         logger.info(String.format("RemoteBulkWriter done! output remote files: %s", getBatchFiles()));
     }
 

sdk-bulkwriter/src/main/java/io/milvus/bulkwriter/RemoteBulkWriter.java

@@ -201,6 +201,14 @@ private static void fillStructType(Types.MessageTypeBuilder messageTypeBuilder,
                 case FloatVector:
                     setMessageType(groupBuilder, PrimitiveType.PrimitiveTypeName.FLOAT, null, subField, true);
                     break;
+                case BinaryVector:
+                case Float16Vector:
+                case BFloat16Vector:
+                case Int8Vector:
+                    boolean isSigned = (subField.getDataType() == io.milvus.v2.common.DataType.Int8Vector);
+                    setMessageType(groupBuilder, PrimitiveType.PrimitiveTypeName.INT32,
+                            LogicalTypeAnnotation.IntLogicalTypeAnnotation.intType(8, isSigned), subField, true);
+                    break;
                 default:
                     break;
             }

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

@@ -0,0 +1,68 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package io.milvus.bulkwriter.common.utils;
+
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+
+public class WriterUtils {
+    private WriterUtils() {
+    }
+
+    public static Object normalizeValue(Object value) {
+        if (value instanceof ByteBuffer) {
+            ByteBuffer byteBuffer = (ByteBuffer) value;
+            byte[] bytes = byteBuffer.array();
+            List<Integer> result = new ArrayList<>(bytes.length);
+            for (byte b : bytes) {
+                result.add((int) b);
+            }
+            return result;
+        }
+        if (value instanceof List) {
+            List<?> values = (List<?>) value;
+            List<Object> normalized = null;
+            for (int i = 0; i < values.size(); i++) {
+                Object item = values.get(i);
+                Object normalizedItem = normalizeValue(item);
+                if (normalized != null) {
+                    normalized.add(normalizedItem);
+                    continue;
+                }
+                if (normalizedItem != item) {
+                    normalized = new ArrayList<>(values.size());
+                    for (int j = 0; j < i; j++) {
+                        normalized.add(values.get(j));
+                    }
+                    normalized.add(normalizedItem);
+                }
+            }
+            return normalized != null ? normalized : value;
+        }
+        if (value instanceof Map) {
+            Map<Object, Object> values = (Map<Object, Object>) value;
+            values.replaceAll((k, v) -> normalizeValue(v));
+            return values;
+        }
+        return value;
+    }
+}

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

@@ -58,6 +58,8 @@
 import java.security.NoSuchAlgorithmException;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.TimeUnit;
 
 import static com.amazonaws.services.s3.internal.Constants.MB;
 
@@ -167,8 +169,18 @@ public void close() {
         if (!closeHttpClient || httpClient == null) {
             return;
         }
-        httpClient.dispatcher().executorService().shutdown();
+        ExecutorService executorService = httpClient.dispatcher().executorService();
+        executorService.shutdown();
         httpClient.connectionPool().evictAll();
+        try {
+            if (!executorService.awaitTermination(5, TimeUnit.SECONDS)) {
+                executorService.shutdownNow();
+            }
+        } catch (InterruptedException e) {
+            executorService.shutdownNow();
+            Thread.currentThread().interrupt();
+            logger.warn("Interrupted while shutting down MinIO HTTP client executor", e);
+        }
         if (httpClient.cache() != null) {
             try {
                 httpClient.cache().close();

sdk-bulkwriter/src/main/java/io/milvus/bulkwriter/storage/client/MinioStorageClient.java

@@ -3,15 +3,14 @@
 import com.google.common.collect.Lists;
 import com.google.gson.Gson;
 import com.google.gson.GsonBuilder;
+import io.milvus.bulkwriter.common.utils.WriterUtils;
 import io.milvus.v2.service.collection.request.CreateCollectionReq;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import java.io.BufferedWriter;
 import java.io.IOException;
-import java.nio.ByteBuffer;
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
 
@@ -43,6 +42,7 @@ private void initWriter() throws IOException {
     @Override
     public void appendRow(Map<String, Object> rowValues, boolean firstWrite) throws IOException {
         rowValues.keySet().removeIf(key -> key.equals(DYNAMIC_FIELD_NAME) && !this.collectionSchema.isEnableDynamicField());
+        rowValues.replaceAll((key, value) -> WriterUtils.normalizeValue(value));
 
         Gson gson = new GsonBuilder().serializeNulls().create();
         List<String> fieldNameList = Lists.newArrayList(rowValues.keySet());
@@ -62,8 +62,6 @@ public void appendRow(Map<String, Object> rowValues, boolean firstWrite) throws
                 String strVal = "";
                 if (val == null) {
                     strVal = nullKey;
-                } else if (val instanceof ByteBuffer) {
-                    strVal = Arrays.toString(((ByteBuffer) val).array());
                 } else if (val instanceof List || val instanceof Map) {
                     strVal = gson.toJson(val); // server-side is using json to parse array field and vector field
                 } else {

sdk-bulkwriter/src/main/java/io/milvus/bulkwriter/writer/CSVFileWriter.java

@@ -2,13 +2,13 @@
 
 import com.google.gson.Gson;
 import com.google.gson.GsonBuilder;
+import io.milvus.bulkwriter.common.utils.WriterUtils;
 import io.milvus.v2.service.collection.request.CreateCollectionReq;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import java.io.BufferedWriter;
 import java.io.IOException;
-import java.nio.ByteBuffer;
 import java.util.Map;
 
 import static io.milvus.param.Constant.DYNAMIC_FIELD_NAME;
@@ -38,7 +38,7 @@ private void initWriter() throws IOException {
     public void appendRow(Map<String, Object> rowValues, boolean firstWrite) throws IOException {
         Gson gson = new GsonBuilder().serializeNulls().create();
         rowValues.keySet().removeIf(key -> key.equals(DYNAMIC_FIELD_NAME) && !this.collectionSchema.isEnableDynamicField());
-        rowValues.replaceAll((key, value) -> value instanceof ByteBuffer ? ((ByteBuffer) value).array() : value);
+        rowValues.replaceAll((key, value) -> WriterUtils.normalizeValue(value));
 
         try {
             if (firstWrite) {