Browse Source

Simplify `createCollection`

jianghua 4 years ago
parent
commit
a86631b22e

+ 91 - 92
src/main/java/io/milvus/client/CollectionMapping.java

@@ -19,119 +19,118 @@
 
 package io.milvus.client;
 
-import java.util.ArrayList;
+import com.google.common.collect.ImmutableMap;
+import io.milvus.grpc.FieldParam;
+import io.milvus.grpc.KeyValuePair;
+import io.milvus.grpc.Mapping;
+import org.json.JSONObject;
+
 import java.util.List;
 import java.util.Map;
-import javax.annotation.Nonnull;
+import java.util.stream.Collectors;
 
 /** Represents a collection mapping */
 public class CollectionMapping {
-  private final String collectionName;
-  private final List<? extends Map<String, Object>> fields;
-  private final String paramsInJson;
-
-  private CollectionMapping(@Nonnull Builder builder) {
-    collectionName = builder.collectionName;
-    fields = builder.fields;
-    paramsInJson = builder.paramsInJson;
+  private final Mapping.Builder builder;
+
+  public static CollectionMapping create(String collectionName) {
+    return new CollectionMapping(collectionName);
   }
 
-  public String getCollectionName() {
-    return collectionName;
+  CollectionMapping(Mapping mapping) {
+    this.builder = mapping.toBuilder();
   }
 
-  public List<? extends Map<String, Object>> getFields() {
-    return fields;
+  private CollectionMapping(String collectionName) {
+    this.builder = Mapping.newBuilder();
+    builder.setCollectionName(collectionName);
   }
 
-  public String getParamsInJson() {
-    return paramsInJson;
+  /**
+   * add a scalar field
+   *
+   * @param name the field name
+   * @param type the field data type
+   * @return this CollectionMapping
+   */
+  public CollectionMapping addField(String name, DataType type) {
+    builder.addFields(FieldParam.newBuilder().setName(name).setTypeValue(type.getVal()).build());
+    return this;
   }
 
-  @Override
-  public String toString() {
-    return String.format(
-        "CollectionMapping = {collectionName = %s, fields = %s, params = %s}",
-        collectionName, fields.toString(), paramsInJson);
+  /**
+   * add a vector field
+   *
+   * @param name the field name
+   * @param type the field data type
+   * @param dimension the vector dimension
+   * @return this CollectionMapping
+   */
+  public CollectionMapping addVectorField(String name, DataType type, int dimension) {
+    FieldParam field = FieldParam.newBuilder()
+        .setName(name)
+        .setTypeValue(type.getVal())
+        .addExtraParams(KeyValuePair.newBuilder()
+            .setKey(MilvusClient.extraParamKey)
+            .setValue(new JSONObject().put("dim", dimension).toString())
+            .build())
+        .build();
+    builder.addFields(field);
+    return this;
   }
 
-  /** Builder for <code>CollectionMapping</code> */
-  public static class Builder {
-    // Required parameters
-    private final String collectionName;
+  public List<Map<String, Object>> getFields() {
+    return builder.getFieldsList().stream().map(f -> ImmutableMap.of(
+        "field", f.getName(),
+        "type", DataType.valueOf(f.getType().getNumber()),
+        "params", (Object) getParamsInJson(f.getExtraParamsList())))
+        .collect(Collectors.toList());
+  }
 
-    // Optional parameters. Default to empty.
-    private List<Map<String, Object>> fields = new ArrayList<>();
-    private String paramsInJson = "{}";
+  /**
+   * Set extra params in json string
+   *
+   * @param paramsInJson Two optional parameters can be included. "segment_row_limit" is default
+   *                     to 100,000. Merge will be triggered if more than this number of entities
+   *                     are inserted into collection. "auto_id" is default to <code>true</code>.
+   *                     Entity ids will be auto-generated by Milvus if set to true.
+   * @return this CollectionMapping
+   */
+  public CollectionMapping setParamsInJson(String paramsInJson) {
+    builder.addExtraParams(KeyValuePair.newBuilder()
+        .setKey(MilvusClient.extraParamKey)
+        .setValue(paramsInJson)
+        .build());
+    return this;
+  }
 
-    /**
-     * @param collectionName collection name
-     */
-    public Builder(@Nonnull String collectionName) {
-      this.collectionName = collectionName;
-    }
+  public String getParamsInJson() {
+    return getParamsInJson(builder.getExtraParamsList());
+  }
 
-    /**
-     * Build with fields. Example fields:
-     * <pre>
-     *   <code>
-     *  [
-     *      {"field": "A", "type": DataType.INT64},
-     *      {"field": "B", "type": DataType.INT32},
-     *      {"field": "C", "type": DataType.FLOAT},
-     *      {"field": "Vec", "type": DataType.VECTOR_FLOAT, "params": {"dim": 128}}
-     *  ]
-     *   </code>
-     * </pre>
-     *
-     * @param fields a list of hashmap listing each field. A field must have keys "field" and
-     *               "type". A vector field must have "dim" in "params". "params" should be in json
-     *               format.
-     * @return <code>Builder</code>
-     */
-    public Builder withFields(@Nonnull List<Map<String, Object>> fields) {
-      this.fields = fields;
-      return this;
-    }
+  public String getCollectionName() {
+    return builder.getCollectionName();
+  }
 
-    /**
-     * Add a single field to collection. Example field:
-     * <pre>
-     *   <code>
-     *      {"field": "A", "type": DataType.INT64}, or
-     *      {"field": "B", "type": DataType.INT32}, or
-     *      {"field": "C", "type": DataType.FLOAT}, or
-     *      {"field": "Vec", "type": DataType.VECTOR_FLOAT, "params": {"dim": 128}}
-     *   </code>
-     * </pre>
-     *
-     * @param field A field must have keys "field" and "type". A vector field must have
-     *              "dim" in "params". <code>FieldBuilder</code> can be used to create a field.
-     * @see FieldBuilder
-     * @return <code>Builder</code>
-     */
-    public Builder field(@Nonnull Map<String, Object> field) {
-      this.fields.add(field);
-      return this;
+  Mapping grpc() {
+    if (builder.getFieldsCount() == 0) {
+      throw new IllegalArgumentException("Fields must not be empty.");
     }
+    return builder.build();
+  }
 
-    /**
-     * Build with extra params in json string format.
-     *
-     * @param paramsInJson Two optional parameters can be included. "segment_row_limit" is default
-     *                     to 100,000. Merge will be triggered if more than this number of entities
-     *                     are inserted into collection. "auto_id" is default to <code>true</code>.
-     *                     Entity ids will be auto-generated by Milvus if set to true.
-     * @see JsonBuilder
-     * @return <code>Builder</code>
-     */
-    public Builder withParamsInJson(String paramsInJson) {
-      this.paramsInJson = paramsInJson;
-      return this;
-    }
+  @Override
+  public String toString() {
+    return String.format(
+        "CollectionMapping = {collectionName = %s, fields = %s, params = %s}",
+        getCollectionName(), getFields(), getParamsInJson());
+  }
 
-    public CollectionMapping build() {
-      return new CollectionMapping(this);
-    }
+  private String getParamsInJson(List<KeyValuePair> extraParams) {
+    return extraParams.stream()
+        .filter(kv -> MilvusClient.extraParamKey.equals(kv.getKey()))
+        .map(KeyValuePair::getValue)
+        .findFirst()
+        .orElse(null);
   }
 }

+ 2 - 13
src/main/java/io/milvus/client/ConnectParam.java

@@ -22,13 +22,12 @@ package io.milvus.client;
 import io.grpc.ManagedChannelBuilder;
 
 import javax.annotation.Nonnull;
+import java.util.Objects;
 import java.util.concurrent.TimeUnit;
 
 /** Contains parameters for connecting to Milvus server */
 public class ConnectParam {
   private final String target;
-  private final String host;
-  private final int port;
   private final String defaultLoadBalancingPolicy;
   private final long connectTimeoutNanos;
   private final long keepAliveTimeNanos;
@@ -37,9 +36,7 @@ public class ConnectParam {
   private final long idleTimeoutNanos;
 
   private ConnectParam(@Nonnull Builder builder) {
-    this.target = builder.target;
-    this.host = builder.host;
-    this.port = builder.port;
+    this.target = builder.target != null ? builder.target : String.format("dns:///%s:%d", builder.host, builder.port);
     this.defaultLoadBalancingPolicy = builder.defaultLoadBalancingPolicy;
     this.connectTimeoutNanos = builder.connectTimeoutNanos;
     this.keepAliveTimeNanos = builder.keepAliveTimeNanos;
@@ -52,14 +49,6 @@ public class ConnectParam {
     return target;
   }
 
-  public String getHost() {
-    return host;
-  }
-
-  public int getPort() {
-    return port;
-  }
-
   public String getDefaultLoadBalancingPolicy() {
     return defaultLoadBalancingPolicy;
   }

+ 5 - 3
src/main/java/io/milvus/client/MilvusClient.java

@@ -32,6 +32,8 @@ import java.util.function.Supplier;
 /** The Milvus Client Interface */
 public interface MilvusClient {
 
+  String extraParamKey = "params";
+
   String clientVersion = new Supplier<String>() {
 
     @Override
@@ -49,6 +51,8 @@ public interface MilvusClient {
     }
   }.get();
 
+  String target();
+
   /** @return current Milvus client version: 0.9.0 */
   default String getClientVersion() {
     return clientVersion;
@@ -83,11 +87,9 @@ public interface MilvusClient {
    * Refer to <code>withFields</code> method for example <code>fields</code> usage.
    * </pre>
    *
-   * @return <code>Response</code>
    * @see CollectionMapping
-   * @see Response
    */
-  Response createCollection(CollectionMapping collectionMapping);
+  void createCollection(CollectionMapping collectionMapping);
 
   /**
    * Checks whether the collection exists

+ 58 - 93
src/main/java/io/milvus/client/MilvusGrpcClient.java

@@ -32,7 +32,10 @@ import io.grpc.ManagedChannel;
 import io.grpc.ManagedChannelBuilder;
 import io.grpc.MethodDescriptor;
 import io.grpc.StatusRuntimeException;
+import io.milvus.client.exception.ClientSideMilvusException;
 import io.milvus.client.exception.InitializationException;
+import io.milvus.client.exception.MilvusException;
+import io.milvus.client.exception.ServerSideMilvusException;
 import io.milvus.client.exception.UnsupportedServerVersion;
 import io.milvus.grpc.*;
 import org.apache.commons.lang3.ArrayUtils;
@@ -53,6 +56,7 @@ import java.util.List;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Function;
+import java.util.function.Supplier;
 
 /** Actual implementation of interface <code>MilvusClient</code> */
 public class MilvusGrpcClient extends AbstractMilvusGrpcClient {
@@ -60,16 +64,16 @@ public class MilvusGrpcClient extends AbstractMilvusGrpcClient {
   private static final Logger logger = LoggerFactory.getLogger(MilvusGrpcClient.class);
   private static final String SUPPORTED_SERVER_VERSION = "0.11";
 
+  private final String target;
   private final ManagedChannel channel;
   private final MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub;
   private final MilvusServiceGrpc.MilvusServiceFutureStub futureStub;
 
   public MilvusGrpcClient(ConnectParam connectParam) {
-    ManagedChannelBuilder builder = connectParam.getTarget() != null
-        ? ManagedChannelBuilder.forTarget(connectParam.getTarget())
-        : ManagedChannelBuilder.forAddress(connectParam.getHost(), connectParam.getPort());
-
-    channel = builder.usePlaintext()
+    target = connectParam.getTarget();
+    channel = ManagedChannelBuilder
+        .forTarget(connectParam.getTarget())
+        .usePlaintext()
         .maxInboundMessageSize(Integer.MAX_VALUE)
         .defaultLoadBalancingPolicy(connectParam.getDefaultLoadBalancingPolicy())
         .keepAliveTime(connectParam.getKeepAliveTime(TimeUnit.NANOSECONDS), TimeUnit.NANOSECONDS)
@@ -84,10 +88,10 @@ public class MilvusGrpcClient extends AbstractMilvusGrpcClient {
       if (response.ok()) {
         String serverVersion = response.getMessage();
         if (!serverVersion.matches("^" + SUPPORTED_SERVER_VERSION + "(\\..*)?$")) {
-          throw new UnsupportedServerVersion(connectParam.getHost(), SUPPORTED_SERVER_VERSION, serverVersion);
+          throw new UnsupportedServerVersion(connectParam.getTarget(), SUPPORTED_SERVER_VERSION, serverVersion);
         }
       } else {
-        throw new InitializationException(connectParam.getHost(), response.getMessage());
+        throw new InitializationException(connectParam.getTarget(), response.getMessage());
       }
     } catch (Throwable t) {
       channel.shutdownNow();
@@ -95,6 +99,11 @@ public class MilvusGrpcClient extends AbstractMilvusGrpcClient {
     }
   }
 
+  @Override
+  public String target() {
+    return target;
+  }
+
   @Override
   protected MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub() {
     return blockingStub;
@@ -138,6 +147,10 @@ public class MilvusGrpcClient extends AbstractMilvusGrpcClient {
         this.futureStub.withInterceptors(timeoutInterceptor);
 
     return new AbstractMilvusGrpcClient() {
+      @Override
+      public String target() {
+        return MilvusGrpcClient.this.target();
+      }
 
       @Override
       protected MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub() {
@@ -184,81 +197,55 @@ public class MilvusGrpcClient extends AbstractMilvusGrpcClient {
 abstract class AbstractMilvusGrpcClient implements MilvusClient {
   private static final Logger logger = LoggerFactory.getLogger(AbstractMilvusGrpcClient.class);
 
-  private final String extraParamKey = "params";
-
   protected abstract MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub();
 
   protected abstract MilvusServiceGrpc.MilvusServiceFutureStub futureStub();
 
   protected abstract boolean maybeAvailable();
 
-  @Override
-  public Response createCollection(@Nonnull CollectionMapping collectionMapping) {
-
-    if (!maybeAvailable()) {
-      logWarning("You are not connected to Milvus server");
-      return new Response(Response.Status.CLIENT_NOT_CONNECTED);
-    }
+  private void translateExceptions(Runnable body) {
+    translateExceptions(() -> {
+      body.run();
+      return null;
+    });
+  }
 
-    List<FieldParam> fields = new ArrayList<>();
-    if (collectionMapping.getFields().size() == 0) {
-      logError("Param fields must not be empty.");
-      return new Response(Response.Status.ILLEGAL_ARGUMENT);
-    }
-    for (Map<String, Object> map : collectionMapping.getFields()) {
-      if (!map.containsKey("field") || !(map.get("field") instanceof String)) {
-        logError("Param fields must contain key 'field' of String.");
-        return new Response(Response.Status.ILLEGAL_ARGUMENT);
-      }
-      if (!map.containsKey("type") || !(map.get("type") instanceof DataType)) {
-        logError("Param fields must contain key 'type' of DataType.");
-        return new Response(Response.Status.ILLEGAL_ARGUMENT);
-      }
-      io.milvus.grpc.FieldParam.Builder fieldParamBuilder = FieldParam.newBuilder()
-          .setName(map.get("field").toString())
-          .setTypeValue(((DataType) map.get("type")).getVal());
-      if (map.containsKey(extraParamKey)) {
-        KeyValuePair extraFieldParam = KeyValuePair.newBuilder()
-            .setKey(extraParamKey)
-            .setValue(map.get(extraParamKey).toString())
-            .build();
-        fieldParamBuilder.addExtraParams(extraFieldParam);
+  @SuppressWarnings("unchecked")
+  private <T> T translateExceptions(Supplier<T> body) {
+    try {
+      T result = body.get();
+      if (result instanceof ListenableFuture) {
+        ListenableFuture futureResult = (ListenableFuture) result;
+        result = (T) Futures.catching(
+            futureResult, Throwable.class, this::translate, MoreExecutors.directExecutor());
       }
-      fields.add(fieldParamBuilder.build());
+      return result;
+    } catch (Throwable e) {
+      return translate(e);
     }
+  }
 
-    Mapping request =
-        Mapping.newBuilder()
-            .setCollectionName(collectionMapping.getCollectionName())
-            .addAllFields(fields)
-            .addExtraParams(KeyValuePair.newBuilder()
-                .setKey(extraParamKey)
-                .setValue(collectionMapping.getParamsInJson())
-                .build())
-            .build();
-
-    Status response;
-
-    try {
-      response = blockingStub().createCollection(request);
+  private <R> R translate(Throwable e) {
+    if (e instanceof MilvusException) {
+      throw (MilvusException) e;
+    } else {
+      throw new ClientSideMilvusException(target(), e);
+    }
+  }
 
-      if (response.getErrorCode() == ErrorCode.SUCCESS) {
-        logInfo("Created collection successfully!\n{}", collectionMapping.toString());
-        return new Response(Response.Status.SUCCESS);
-      } else if (response.getReason().contentEquals("Collection already exists")) {
-        logWarning("Collection `{}` already exists", collectionMapping.getCollectionName());
-        return new Response(
-            Response.Status.valueOf(response.getErrorCodeValue()), response.getReason());
-      } else {
-        logError(
-            "Create collection failed\n{}\n{}", collectionMapping.toString(), response.toString());
-        return new Response(
-            Response.Status.valueOf(response.getErrorCodeValue()), response.getReason());
-      }
-    } catch (StatusRuntimeException e) {
-      logError("createCollection RPC failed:\n{}", e.getStatus().toString());
-      return new Response(Response.Status.RPC_ERROR, e.toString());
+  private Void checkResponseStatus(Status status) {
+    if (status.getErrorCode() != ErrorCode.SUCCESS) {
+      throw new ServerSideMilvusException(target(), status);
     }
+    return null;
+  }
+
+  @Override
+  public void createCollection(@Nonnull CollectionMapping collectionMapping) {
+    translateExceptions(() -> {
+      Status response = blockingStub().createCollection(collectionMapping.grpc());
+      checkResponseStatus(response);
+    });
   }
 
   @Override
@@ -1130,29 +1117,7 @@ abstract class AbstractMilvusGrpcClient implements MilvusClient {
       response = blockingStub().describeCollection(request);
 
       if (response.getStatus().getErrorCode() == ErrorCode.SUCCESS) {
-        String extraParam = "";
-        for (KeyValuePair kv : response.getExtraParamsList()) {
-          if (kv.getKey().contentEquals(extraParamKey)) {
-            extraParam = kv.getValue();
-          }
-        }
-        // convert fields to list of hashmap
-        List<FieldParam> fields = response.getFieldsList();
-        List<Map<String, Object>> fieldsCollection = new ArrayList<>(fields.size());
-        for (FieldParam fieldParam : fields) {
-          Map<String, Object> map = new HashMap<>();
-          // copy from fieldParam to map
-          map.put("field", fieldParam.getName());
-          map.put("type", fieldParam.getType());
-          map.put("indexParams", kvListToString(fieldParam.getIndexParamsList()));
-          map.put("params", kvListToString(fieldParam.getExtraParamsList()));
-          fieldsCollection.add(map);
-        }
-        CollectionMapping collectionMapping =
-            new CollectionMapping.Builder(response.getCollectionName())
-                .withFields(fieldsCollection)
-                .withParamsInJson(extraParam)
-                .build();
+        CollectionMapping collectionMapping = new CollectionMapping(response);
         logInfo("Get Collection Info `{}` returned:\n{}", collectionName, collectionMapping);
         return new GetCollectionInfoResponse(
             new Response(Response.Status.SUCCESS), collectionMapping);

+ 16 - 0
src/main/java/io/milvus/client/exception/ClientSideMilvusException.java

@@ -0,0 +1,16 @@
+package io.milvus.client.exception;
+
+public class ClientSideMilvusException extends MilvusException {
+
+  public ClientSideMilvusException(String target) {
+    super(target, false, null, null);
+  }
+
+  public ClientSideMilvusException(String target, Throwable cause) {
+    super(target, false, null, cause);
+  }
+
+  public ClientSideMilvusException(String target, String message) {
+    super(target, false, message, null);
+  }
+}

+ 3 - 3
src/main/java/io/milvus/client/exception/InitializationException.java

@@ -1,7 +1,7 @@
 package io.milvus.client.exception;
 
-public class InitializationException extends MilvusException {
-  public InitializationException(String host, String message) {
-    super(false, host + ": " + message);
+public class InitializationException extends ClientSideMilvusException {
+  public InitializationException(String target, String message) {
+    super(target, message);
   }
 }

+ 14 - 8
src/main/java/io/milvus/client/exception/MilvusException.java

@@ -1,20 +1,26 @@
 package io.milvus.client.exception;
 
-public class MilvusException extends RuntimeException {
+public abstract class MilvusException extends RuntimeException {
+  private String target;
   private boolean fillInStackTrace;
 
-  MilvusException(boolean fillInStackTrace) {
-    this.fillInStackTrace = fillInStackTrace;
+  MilvusException(String target, boolean fillInStackTrace) {
+    this(target, fillInStackTrace, null, null);
   }
 
-  MilvusException(boolean fillInStackTrace, Throwable cause) {
-    super(cause);
+  MilvusException(String target, boolean fillInStackTrace, String message, Throwable cause) {
+    super(message, cause);
+    this.target = target;
     this.fillInStackTrace = fillInStackTrace;
   }
 
-  MilvusException(boolean fillInStackTrace, String message) {
-    super(message);
-    this.fillInStackTrace = fillInStackTrace;
+  @Override
+  public final String getMessage() {
+    return String.format("%s: %s", target, getErrorMessage());
+  }
+
+  protected String getErrorMessage() {
+    return super.getMessage();
   }
 
   @Override

+ 33 - 0
src/main/java/io/milvus/client/exception/ServerSideMilvusException.java

@@ -0,0 +1,33 @@
+package io.milvus.client.exception;
+
+import io.milvus.grpc.ErrorCode;
+import io.milvus.grpc.Status;
+
+public class ServerSideMilvusException extends MilvusException {
+  private ErrorCode errorCode;
+  private String reason;
+
+  public ServerSideMilvusException(String target, Status status) {
+    super(target, false);
+    this.errorCode = status.getErrorCode();
+    this.reason = status.getReason();
+  }
+
+  public ErrorCode getErrorCode() {
+    return errorCode;
+  }
+
+  public String getReason() {
+    return reason;
+  }
+
+  @Override
+  public synchronized Throwable fillInStackTrace() {
+    return this;
+  }
+
+  @Override
+  public String getErrorMessage() {
+    return String.format("ServerSideMilvusException{errorCode=%s, reason=%s}", errorCode, reason);
+  }
+}

+ 6 - 8
src/main/java/io/milvus/client/exception/UnsupportedServerVersion.java

@@ -2,21 +2,19 @@ package io.milvus.client.exception;
 
 import io.milvus.client.MilvusClient;
 
-public class UnsupportedServerVersion extends MilvusException {
-  private String host;
+public class UnsupportedServerVersion extends ClientSideMilvusException {
   private String expect;
   private String actual;
 
-  public UnsupportedServerVersion(String host, String expect, String actual) {
-    super(false);
-    this.host = host;
+  public UnsupportedServerVersion(String target, String expect, String actual) {
+    super(target);
     this.expect = expect;
     this.actual = actual;
   }
 
   @Override
-  public String getMessage() {
-    return String.format("%s: Milvus client %s is expected to work with Milvus server %s, but the version of the connected server is %s",
-        host, MilvusClient.clientVersion, expect, actual);
+  public String getErrorMessage() {
+    return String.format("Milvus client %s is expected to work with Milvus server %s, but the version of the connected server is %s",
+        MilvusClient.clientVersion, expect, actual);
   }
 }

+ 55 - 68
src/test/java/io/milvus/client/MilvusGrpcClientTest.java

@@ -26,9 +26,11 @@ import com.google.common.util.concurrent.MoreExecutors;
 import io.grpc.NameResolverProvider;
 import io.grpc.NameResolverRegistry;
 import io.milvus.client.InsertParam.Builder;
-import io.milvus.client.Response.Status;
+import io.milvus.client.exception.ClientSideMilvusException;
 import io.milvus.client.exception.InitializationException;
+import io.milvus.client.exception.ServerSideMilvusException;
 import io.milvus.client.exception.UnsupportedServerVersion;
+import io.milvus.grpc.ErrorCode;
 import org.apache.commons.lang3.ArrayUtils;
 import org.apache.commons.text.RandomStringGenerator;
 import org.checkerframework.checker.nullness.compatqual.NullableDecl;
@@ -123,6 +125,10 @@ class MilvusClientTest {
     return new ConnectParam.Builder().withHost(host).withPort(port);
   }
 
+  protected void assertErrorCode(ErrorCode errorCode, Runnable runnable) {
+    assertEquals(errorCode, assertThrows(ServerSideMilvusException.class, runnable::run).getErrorCode());
+  }
+
   // Helper function that generates random float vectors
   static List<List<Float>> generateFloatVectors(int vectorCount, int dimension) {
     SplittableRandom splittableRandom = new SplittableRandom();
@@ -207,19 +213,18 @@ class MilvusClientTest {
     randomCollectionName = generator.generate(10);
     size = 100000;
     dimension = 128;
-    CollectionMapping collectionMapping =
-        new CollectionMapping.Builder(randomCollectionName)
-            .field(new FieldBuilder("int64", DataType.INT64).build())
-            .field(new FieldBuilder("float", DataType.FLOAT).build())
-            .field(new FieldBuilder("float_vec", DataType.VECTOR_FLOAT)
-                .param("dim", dimension)
-                .build())
-            .withParamsInJson(new JsonBuilder().param("segment_row_limit", 50000)
-                                               .param("auto_id", false)
-                                               .build())
-            .build();
 
-    assertTrue(client.createCollection(collectionMapping).ok());
+    CollectionMapping collectionMapping = CollectionMapping
+        .create(randomCollectionName)
+        .addField("int64", DataType.INT64)
+        .addField("float", DataType.FLOAT)
+        .addVectorField("float_vec", DataType.VECTOR_FLOAT, dimension)
+        .setParamsInJson(new JsonBuilder()
+            .param("segment_row_limit", 50000)
+            .param("auto_id", false)
+            .build());
+
+    client.createCollection(collectionMapping);
   }
 
   @org.junit.jupiter.api.AfterEach
@@ -310,37 +315,25 @@ class MilvusClientTest {
   @org.junit.jupiter.api.Test
   void createInvalidCollection() {
     // invalid collection name
-    String invalidCollectionName = "╯°□°)╯";
-    CollectionMapping invalidCollectionMapping =
-        new CollectionMapping.Builder(invalidCollectionName)
-            .field(new FieldBuilder("float_vec", DataType.VECTOR_FLOAT)
-                .param("dim", dimension)
-                .build())
-            .build();
-    Response createCollectionResponse = client.createCollection(invalidCollectionMapping);
-    assertFalse(createCollectionResponse.ok());
-    assertEquals(Response.Status.ILLEGAL_COLLECTION_NAME, createCollectionResponse.getStatus());
+    CollectionMapping invalidCollectionName = CollectionMapping
+        .create("╯°□°)╯")
+        .addVectorField("float_vec", DataType.VECTOR_FLOAT, dimension);
+
+    assertErrorCode(ErrorCode.ILLEGAL_COLLECTION_NAME, () -> client.createCollection(invalidCollectionName));
 
     // invalid field
-    invalidCollectionMapping =
-        new CollectionMapping.Builder("validCollectionName")
-            .build();
-    createCollectionResponse = client.createCollection(invalidCollectionMapping);
-    assertFalse(createCollectionResponse.ok());
-
-    // invalid segment_row_limit
-    invalidCollectionMapping =
-        new CollectionMapping.Builder("validCollectionName")
-            .field(new FieldBuilder("int64", DataType.INT64).build())
-            .field(new FieldBuilder("float", DataType.FLOAT).build())
-            .field(new FieldBuilder("float_vec", DataType.VECTOR_FLOAT)
-                .param("dim", dimension)
-                .build())
-            .withParamsInJson(new JsonBuilder().param("segment_row_limit", -1000).build())
-            .build();
-    createCollectionResponse = client.createCollection(invalidCollectionMapping);
-    assertFalse(createCollectionResponse.ok());
-    assertEquals(Status.ILLEGAL_ARGUMENT, createCollectionResponse.getStatus());
+    CollectionMapping withoutField = CollectionMapping.create("validCollectionName");
+    assertThrows(ClientSideMilvusException.class, () -> client.createCollection(withoutField));
+
+    // invalid segment_row_count
+    CollectionMapping invalidSegmentRowCount = CollectionMapping
+        .create("validCollectionName")
+        .addField("int64", DataType.INT64)
+        .addField("float", DataType.FLOAT)
+        .addVectorField("float_vec", DataType.VECTOR_FLOAT, dimension)
+        .setParamsInJson(new JsonBuilder().param("segment_row_limit", -1000).build());
+
+    assertErrorCode(ErrorCode.ILLEGAL_ARGUMENT, () -> client.createCollection(invalidSegmentRowCount));
   }
 
   @org.junit.jupiter.api.Test
@@ -593,13 +586,12 @@ class MilvusClientTest {
     final int binaryDimension = 10000;
 
     String binaryCollectionName = generator.generate(10);
-    CollectionMapping collectionMapping =
-        new CollectionMapping.Builder(binaryCollectionName)
-            .field(new FieldBuilder("binary_vec", DataType.VECTOR_BINARY)
-                .param("dim", binaryDimension)
-                .build())
-            .build();
-    assertTrue(client.createCollection(collectionMapping).ok());
+
+    CollectionMapping collectionMapping = CollectionMapping
+        .create(binaryCollectionName)
+        .addVectorField("binary_vec", DataType.VECTOR_BINARY, binaryDimension);
+
+    client.createCollection(collectionMapping);
 
     List<List<Byte>> vectors = generateBinaryVectors(size, binaryDimension);
     InsertParam insertParam =
@@ -758,15 +750,13 @@ class MilvusClientTest {
     final int binaryDimension = 64;
 
     String binaryCollectionName = generator.generate(10);
-    CollectionMapping collectionMapping =
-        new CollectionMapping.Builder(binaryCollectionName)
-            .field(new FieldBuilder("int64", DataType.INT64).build())
-            .field(new FieldBuilder("float", DataType.FLOAT).build())
-            .field(new FieldBuilder("binary_vec", DataType.VECTOR_BINARY)
-                .param("dim", binaryDimension)
-                .build())
-            .build();
-    assertTrue(client.createCollection(collectionMapping).ok());
+    CollectionMapping collectionMapping = CollectionMapping
+        .create(binaryCollectionName)
+        .addField("int64", DataType.INT64)
+        .addField("float", DataType.FLOAT)
+        .addVectorField("binary_vec", DataType.VECTOR_BINARY, binaryDimension);
+
+    client.createCollection(collectionMapping);
 
     // field list for insert
     List<Long> intValues = new ArrayList<>(size);
@@ -836,8 +826,7 @@ class MilvusClientTest {
         .get().getFields();
     for (Map<String, Object> field : fields) {
       if (field.get("field").equals("float_vec")) {
-        JSONObject jsonObject = new JSONObject(field.get("params").toString());
-        JSONObject params = new JSONObject(jsonObject.get("params").toString());
+        JSONObject params = new JSONObject(field.get("params").toString());
         assertTrue(params.has("dim"));
       }
     }
@@ -955,14 +944,12 @@ class MilvusClientTest {
     final int binaryDimension = 64;
 
     String binaryCollectionName = generator.generate(10);
-    CollectionMapping collectionMapping =
-        new CollectionMapping.Builder(binaryCollectionName)
-            .field(new FieldBuilder("binary_vec", DataType.VECTOR_BINARY)
-                .param("dim", binaryDimension)
-                .build())
-            .withParamsInJson(new JsonBuilder().param("auto_id", false).build())
-            .build();
-    assertTrue(client.createCollection(collectionMapping).ok());
+    CollectionMapping collectionMapping = CollectionMapping
+        .create(binaryCollectionName)
+        .addVectorField("binary_vec", DataType.VECTOR_BINARY, binaryDimension)
+        .setParamsInJson(new JsonBuilder().param("auto_id", false).build());
+
+    client.createCollection(collectionMapping);
 
     List<List<Byte>> vectors = generateBinaryVectors(size, binaryDimension);
     List<Long> entityIds = LongStream.range(0, size).boxed().collect(Collectors.toList());