Updates to Milvus adapter

nosqlbench · Apr 15, 2024 · 41911a0 · 41911a0
1 parent af3928b
commit 41911a0
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 19 deletions.
diff --git a/nb-adapters/adapter-milvus/pom.xml b/nb-adapters/adapter-milvus/pom.xml
@@ -50,12 +50,12 @@
         <dependency>
             <groupId>com.google.protobuf</groupId>
             <artifactId>protobuf-java-util</artifactId>
-            <version>3.24.0</version>
+            <version>3.25.3</version>
         </dependency>
         <dependency>
             <groupId>io.milvus</groupId>
             <artifactId>milvus-sdk-java</artifactId>
-            <version>2.3.4</version>
+            <version>2.3.5</version>
         </dependency>
 
     </dependencies>

diff --git a/nb-adapters/adapter-milvus/src/main/java/io/nosqlbench/adapter/milvus/MilvusSpace.java b/nb-adapters/adapter-milvus/src/main/java/io/nosqlbench/adapter/milvus/MilvusSpace.java
@@ -51,10 +51,8 @@ public class MilvusSpace implements AutoCloseable {
      * Create a new MilvusSpace Object which stores all stateful contextual information needed to interact
      * with the Milvus/Zilliz database instance.
      *
-     * @param name
-     *     The name of this space
-     * @param cfg
-     *     The configuration ({@link NBConfiguration}) for this nb run
+     * @param name The name of this space
+     * @param cfg  The configuration ({@link NBConfiguration}) for this nb run
      */
     public MilvusSpace(String name, NBConfiguration cfg) {
         this.name = name;
@@ -89,16 +87,13 @@ private MilvusServiceClient createClient() {
             ).orElseGet(
                 () -> cfg.getOptional("token")
                     .orElseThrow(() -> new RuntimeException("You must provide either a token_file or a token to " +
-                        "configure a Milvus client"))
+                        "configure a Milvus/Zilliz client"))
             );
-
         builder = builder.withToken(requiredToken);
 
         ConnectParam connectParams = builder.build();
-
-        logger.info(this.name + ": Creating new Milvus/Zilliz Client with (masked) " +
-            "token [" + MilvusAdapterUtils.maskDigits(builder.getToken()) + "], uri/endpoint [" + builder.getUri() + "]"
-        );
+        logger.info("{}: Creating new Milvus/Zilliz Client with (masked) token [{}], uri/endpoint [{}]",
+            this.name, MilvusAdapterUtils.maskDigits(builder.getToken()), builder.getUri());
         return new MilvusServiceClient(connectParams);
     }
 
@@ -117,8 +112,8 @@ public static NBConfigModel getConfigModel() {
                     .setDescription("the URI endpoint in which the database is running.")
             )
             .add(
-                Param.optional(List.of("database_name","database"))
-                    .setDescription("the name of the database to use. Defaults to 'baselines'")
+                Param.optional(List.of("database_name", "database"))
+                    .setDescription("the name of the database to use.")
             )
             .asReadOnly();
     }

diff --git a/...ain/java/io/nosqlbench/adapter/milvus/opdispensers/MilvusCreateCollectionOpDispenser.java b/...ain/java/io/nosqlbench/adapter/milvus/opdispensers/MilvusCreateCollectionOpDispenser.java
@@ -19,6 +19,7 @@
 import io.milvus.client.MilvusServiceClient;
 import io.milvus.common.clientenum.ConsistencyLevelEnum;
 import io.milvus.grpc.DataType;
+import io.milvus.param.collection.CollectionSchemaParam;
 import io.milvus.param.collection.CreateCollectionParam;
 import io.milvus.param.collection.FieldType;
 import io.nosqlbench.adapter.milvus.MilvusDriverAdapter;
@@ -70,11 +71,10 @@ public LongFunction<CreateCollectionParam> getParamFunc(
             CreateCollectionParam.Builder::withDatabaseName);
 
         List<FieldType> fieldTypes = buildFieldTypesStruct(
-            op.getAsSubOps("field_types", ParsedOp.SubOpNaming.SubKey),
-            ebF
+            op.getAsSubOps("field_types", ParsedOp.SubOpNaming.SubKey)
         );
         final LongFunction<CreateCollectionParam.Builder> f = ebF;
-        ebF = l -> f.apply(l).withFieldTypes(fieldTypes);
+        ebF = l -> f.apply(l).withSchema(CollectionSchemaParam.newBuilder().withFieldTypes(fieldTypes).build());
 
         final LongFunction<CreateCollectionParam.Builder> lastF = ebF;
         return l -> lastF.apply(l).build();
@@ -95,10 +95,9 @@ public LongFunction<MilvusBaseOp<CreateCollectionParam>> createOpFunc(
      * Function to build the {@link FieldType}s for the {@link CreateCollectionParam}.
      *
      * @param fieldTypesData The static map of config data from the create collection request
-     * @param ebF
      * @return a list of static field types
      */
-    private List<FieldType> buildFieldTypesStruct(Map<String, ParsedOp> fieldTypesData, LongFunction<CreateCollectionParam.Builder> ebF) {
+    private List<FieldType> buildFieldTypesStruct(Map<String, ParsedOp> fieldTypesData) {
         List<FieldType> fieldTypes = new ArrayList<>();
         fieldTypesData.forEach((name, fieldspec) -> {
             FieldType.Builder builder = FieldType.newBuilder()