Merge pull request #148 from sahuang/0.9.2

Fix MilvusService

examples/src/main/java/MilvusIndexExample.java

@@ -76,18 +76,20 @@ public class MilvusIndexExample {
     ConnectParam connectParam = new ConnectParam.Builder().build();
     MilvusClient client = new MilvusGrpcClient(connectParam);
 
-    final String collectionName = "demo_index";
-    if (client.listCollections().contains(collectionName)) {
-      client.dropCollection(collectionName);
-    }
-
     /*
      * Basic create collection:
-     *   Another way to create a collection is to predefine a schema.
-     *   The schema can later be used in querying.
+     *   Another way to create a collection is to predefine a schema. Then we can use a
+     *   MilvusService to wrap the schema, collection name and Milvus client. The service
+     *   is intended to simplify API calls.
      */
+    final String collectionName = "demo_index";
     FilmSchema filmSchema = new FilmSchema();
     MilvusService service = new MilvusService(client, collectionName, filmSchema);
+
+    if (service.listCollections().contains(collectionName)) {
+      service.dropCollection();
+    }
+
     service.createCollection(
         new JsonBuilder().param("auto_id", false).param("segment_row_limit", 4096).build());
 
@@ -129,16 +131,14 @@ public class MilvusIndexExample {
     csvReader.close();
 
     // Now we can insert entities, the total row count should be 8657.
-    InsertParam insertParam =
-        InsertParam.create(collectionName)
-            .addField("release_year", DataType.INT32, releaseYears)
-            .addVectorField("embedding", DataType.VECTOR_FLOAT, embeddings)
-            .setEntityIds(ids);
-
-    client.insert(insertParam);
-    client.flush(collectionName);
-    System.out.printf(
-        "There are %d films in the collection.\n", client.countEntities(collectionName));
+    service.insert(
+        insertParam ->
+            insertParam
+                .withIds(ids)
+                .with(filmSchema.releaseYear, releaseYears)
+                .with(filmSchema.embedding, embeddings));
+    service.flush();
+    System.out.printf("There are %d films in the collection.\n", service.countEntities());
 
     /*
      * Basic create index:
@@ -149,18 +149,27 @@ public class MilvusIndexExample {
      *
      *   Note that if there is already an index and create index is called again, the previous index
      *   will be replaced.
+     *
+     *   We present two ways to create an index: using Index or MilvusService.
      */
+    // Approach one
     Index index =
         Index.create(collectionName, "embedding")
             .setIndexType(IndexType.IVF_FLAT)
             .setMetricType(MetricType.L2)
             .setParamsInJson(new JsonBuilder().param("nlist", 100).build());
-
     client.createIndex(index);
 
+    // Approach two
+    service.createIndex(
+        filmSchema.embedding,
+        IndexType.IVF_FLAT,
+        MetricType.L2,
+        new JsonBuilder().param("nlist", 100).build());
+
     // Get collection stats with index
     System.out.println("\n--------Collection Stats--------");
-    JSONObject json = new JSONObject(client.getCollectionStats(collectionName));
+    JSONObject json = new JSONObject(service.getCollectionStats());
     System.out.println(json.toString(4));
 
     /*
@@ -175,14 +184,18 @@ public class MilvusIndexExample {
      */
     List<List<Float>> queryEmbedding = randomFloatVectors(1, dimension);
     final int topK = 3;
-    Query query = Query.bool(Query.must(
-        filmSchema.releaseYear.in(1995, 2002),
-        filmSchema.embedding.query(queryEmbedding)
-            .metricType(MetricType.L2)
-            .top(topK)
-            .param("nprobe", 8)));
-    SearchParam searchParam = service.buildSearchParam(query)
-        .setParamsInJson("{\"fields\": [\"release_year\", \"embedding\"]}");
+    Query query =
+        Query.bool(
+            Query.must(
+                filmSchema.releaseYear.in(1995, 2002),
+                filmSchema.embedding.query(queryEmbedding)
+                    .metricType(MetricType.L2)
+                    .top(topK)
+                    .param("nprobe", 8)));
+    SearchParam searchParam =
+        service
+            .buildSearchParam(query)
+            .setParamsInJson("{\"fields\": [\"release_year\", \"embedding\"]}");
     System.out.println("\n--------Search Result--------");
     SearchResult searchResult = service.search(searchParam);
     System.out.println("- ids: " + searchResult.getResultIdsList().toString());
@@ -203,14 +216,14 @@ public class MilvusIndexExample {
      * Basic delete index:
      *   Index can be dropped for a vector field.
      */
-    client.dropIndex(collectionName, "embedding");
+    service.dropIndex(filmSchema.embedding);
 
-    if (client.listCollections().contains(collectionName)) {
-      client.dropCollection(collectionName);
+    if (service.listCollections().contains(collectionName)) {
+      service.dropCollection();
     }
 
     // Close connection
-    client.close();
+    service.close();
   }
 
   // Schema that defines a collection

src/main/java/io/milvus/client/CollectionMapping.java

@@ -47,7 +47,7 @@ public class CollectionMapping {
   }
 
   /**
-   * add a scalar field
+   * Add a scalar field to the collection schema.
    *
    * @param name the field name
    * @param type the field data type
@@ -59,7 +59,7 @@ public class CollectionMapping {
   }
 
   /**
-   * add a vector field
+   * Add a vector field to the collection schema.
    *
    * @param name the field name
    * @param type the field data type
@@ -106,7 +106,7 @@ public class CollectionMapping {
    * 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
+   *     524,288. Merge will be triggered if more than this number of entities are inserted into the
    *     collection. "auto_id" is default to <code>true</code>. Entity ids will be auto-generated by
    *     Milvus if set to true.
    * @return this CollectionMapping

src/main/java/io/milvus/client/DataType.java

@@ -24,17 +24,17 @@ import java.util.Optional;
 
 /** Represents available data types. */
 public enum DataType {
-  NONE(0),
-  BOOL(1),
-  INT8(2),
-  INT16(3),
+  // NONE(0),
+  // BOOL(1),
+  // INT8(2),
+  // INT16(3),
   INT32(4),
   INT64(5),
 
   FLOAT(10),
   DOUBLE(11),
 
-  STRING(20),
+  // STRING(20),
 
   VECTOR_BINARY(100),
   VECTOR_FLOAT(101),

src/main/java/io/milvus/client/Index.java

@@ -28,7 +28,7 @@ import javax.annotation.Nonnull;
 
 /**
  * Represents an index containing <code>fieldName</code>, <code>indexName</code> and <code>
- * paramsInJson</code>, which contains index_type, params etc.
+ * paramsInJson</code>, which contains <code>index_type</code>, params etc.
  */
 public class Index {
   private final IndexParam.Builder builder;
@@ -50,6 +50,7 @@ public class Index {
     return builder.getCollectionName();
   }
 
+  /** @param collectionName The collection name */
   public Index setCollectionName(@Nonnull String collectionName) {
     builder.setCollectionName(collectionName);
     return this;
@@ -59,8 +60,9 @@ public class Index {
     return builder.getFieldName();
   }
 
-  public Index setFieldName(@Nonnull String collectionName) {
-    builder.setFieldName(collectionName);
+  /** @param fieldName The field name */
+  public Index setFieldName(@Nonnull String fieldName) {
+    builder.setFieldName(fieldName);
     return this;
   }
 
@@ -72,10 +74,12 @@ public class Index {
     return toMap(builder.getExtraParamsList());
   }
 
+  /** @param indexType The index type */
   public Index setIndexType(IndexType indexType) {
     return addParam("index_type", indexType.name());
   }
 
+  /** @param metricType The metric type */
   public Index setMetricType(MetricType metricType) {
     return addParam("metric_type", metricType.name());
   }

src/main/java/io/milvus/client/MilvusClient.java

@@ -97,7 +97,7 @@ public interface MilvusClient extends AutoCloseable {
   void createCollection(CollectionMapping collectionMapping);
 
   /**
-   * Checks whether the collection exists
+   * Check whether a collection exists
    *
    * @param collectionName collection to check
    * @return true if the collection exists, false otherwise.

src/main/java/io/milvus/client/SearchParam.java

@@ -49,15 +49,18 @@ public class SearchParam {
     builder.setCollectionName(collectionName);
   }
 
+  /** @param collectionName collection name */
   public static SearchParam create(String collectionName) {
     return new SearchParam(collectionName);
   }
 
+  /** @param json DSL in JSON format. */
   public SearchParam setDsl(JSONObject json) {
     builder.setDsl(json.toString());
     return this;
   }
 
+  /** @param dsl DSL in string format. */
   public SearchParam setDsl(String dsl) {
     try {
       JSONObject dslJson = new JSONObject(dsl);

src/main/java/io/milvus/client/dsl/BoolQuery.java

@@ -6,6 +6,7 @@ import java.util.stream.Collectors;
 import org.json.JSONArray;
 import org.json.JSONObject;
 
+/** Boolean Query */
 public class BoolQuery extends Query {
   private final Type type;
   private final List<Query> subqueries;

src/main/java/io/milvus/client/dsl/MilvusService.java

@@ -2,6 +2,8 @@ package io.milvus.client.dsl;
 
 import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
+import io.milvus.client.CollectionMapping;
+import io.milvus.client.CompactParam;
 import io.milvus.client.Index;
 import io.milvus.client.IndexType;
 import io.milvus.client.MetricType;
@@ -15,6 +17,11 @@ import java.util.concurrent.TimeUnit;
 import java.util.function.Consumer;
 import java.util.stream.Collectors;
 
+/**
+ * A service that wraps client, collection name and schema together to simplify API calls. It is
+ * recommended to use the service if you need a schema for Milvus operations. All operations can
+ * then be called with <code>MilvusService</code> instead of <code>MilvusClient</code>.
+ */
 public class MilvusService {
   private final MilvusClient client;
   private final String collectionName;
@@ -26,26 +33,54 @@ public class MilvusService {
     this.schema = schema;
   }
 
-  public MilvusService withTimeout(int timeout, TimeUnit unit) {
-    return new MilvusService(client.withTimeout(timeout, unit), collectionName, schema);
+  /**
+   * Milvus service with timeout support.
+   *
+   * @param timeout the desired timeout
+   * @param timeoutUnit unit for timeout
+   */
+  public MilvusService withTimeout(int timeout, TimeUnit timeoutUnit) {
+    return new MilvusService(client.withTimeout(timeout, timeoutUnit), collectionName, schema);
   }
 
+  /** Close the client. Wait at most 1 minute for graceful shutdown. */
   public void close() {
     client.close();
   }
 
+  /** Count entities in the current collection */
   public long countEntities() {
     return client.countEntities(collectionName);
   }
 
+  /** Create collection with predefined schema. */
   public void createCollection() {
     createCollection("{}");
   }
 
+  /**
+   * Create collection with predefined schema.
+   *
+   * @param paramsInJson Set optional "segment_row_limit" or "auto_id".
+   */
   public void createCollection(String paramsInJson) {
     client.createCollection(schema.mapToCollection(collectionName).setParamsInJson(paramsInJson));
   }
 
+  /**
+   * Create index with schema field.
+   *
+   * <pre>
+   * example usage:
+   * <code>
+   * service.createIndex(
+   *     schema.embedding,
+   *     IndexType.IVF_FLAT,
+   *     MetricType.L2,
+   *     new JsonBuilder().param("nlist", 100).build());
+   * </code>
+   * </pre>
+   */
   public void createIndex(
       Schema.VectorField vectorField,
       IndexType indexType,
@@ -54,6 +89,20 @@ public class MilvusService {
     Futures.getUnchecked(createIndexAsync(vectorField, indexType, metricType, paramsInJson));
   }
 
+  /**
+   * Create index with schema field.
+   *
+   * <pre>
+   * example usage:
+   * <code>
+   * service.createIndex(
+   *     schema.embedding,
+   *     IndexType.IVF_FLAT,
+   *     MetricType.L2,
+   *     new JsonBuilder().param("nlist", 100).build());
+   * </code>
+   * </pre>
+   */
   public ListenableFuture<Void> createIndexAsync(
       Schema.VectorField vectorField,
       IndexType indexType,
@@ -66,55 +115,225 @@ public class MilvusService {
             .setParamsInJson(paramsInJson));
   }
 
+  /** Delete entities by IDs. */
   public void deleteEntityByID(List<Long> ids) {
     client.deleteEntityByID(collectionName, ids);
   }
 
+  /** Drop the current collection. */
   public void dropCollection() {
     client.dropCollection(collectionName);
   }
 
+  /** Flush the current collection. */
   public void flush() {
     client.flush(collectionName);
   }
 
+  /** Flush the current collection. */
   public ListenableFuture<Void> flushAsync() {
     return client.flushAsync(collectionName);
   }
 
+  /**
+   * Get entity by IDs.
+   *
+   * @param ids The ids to get
+   * @return Map of id to <code>Schema.Entity</code> value.
+   */
   public Map<Long, Schema.Entity> getEntityByID(List<Long> ids) {
     return getEntityByID(ids, Collections.emptyList());
   }
 
+  /**
+   * Get entity by IDs.
+   *
+   * @param ids The ids to get
+   * @param fields The fields to return values for
+   * @return Map of id to <code>Schema.Entity</code> value.
+   */
   public Map<Long, Schema.Entity> getEntityByID(List<Long> ids, List<Schema.Field<?>> fields) {
     List<String> fieldNames = fields.stream().map(f -> f.name).collect(Collectors.toList());
     return client.getEntityByID(collectionName, ids, fieldNames).entrySet().stream()
         .collect(Collectors.toMap(e -> e.getKey(), e -> schema.new Entity(e.getValue())));
   }
 
+  /**
+   * Check whether a collection exists
+   *
+   * @param collectionName collection to check
+   * @return true if the collection exists, false otherwise.
+   */
   public boolean hasCollection(String collectionName) {
     return client.hasCollection(collectionName);
   }
 
+  /**
+   * Insert data with schema.
+   *
+   * <pre>
+   * example usage:
+   * <code>
+   * service.insert(
+   *         insertParam ->
+   *             insertParam
+   *                 .withIds(ids)
+   *                 .with(schema.int_field, int_values)
+   *                 .with(schema.embedding, embeddings));
+   * </code>
+   * </pre>
+   *
+   * @return a list of ids of the inserted entities
+   */
   public List<Long> insert(Consumer<InsertParam> insertParamBuilder) {
     return Futures.getUnchecked(insertAsync(insertParamBuilder));
   }
 
+  /**
+   * Insert data with schema.
+   *
+   * <pre>
+   * example usage:
+   * <code>
+   * service.insert(
+   *         insertParam ->
+   *             insertParam
+   *                 .withIds(ids)
+   *                 .with(schema.int_field, int_values)
+   *                 .with(schema.embedding, embeddings));
+   * </code>
+   * </pre>
+   *
+   * @return a list of ids of the inserted entities
+   */
   public ListenableFuture<List<Long>> insertAsync(Consumer<InsertParam> insertParamBuilder) {
     InsertParam insertParam = schema.insertInto(collectionName);
     insertParamBuilder.accept(insertParam);
     return client.insertAsync(insertParam.getInsertParam());
   }
 
+  /** Search with searchParam. */
   public SearchResult search(SearchParam searchParam) {
     return client.search(searchParam);
   }
 
+  /** Search with searchParam. */
   public ListenableFuture<SearchResult> searchAsync(SearchParam searchParam) {
     return client.searchAsync(searchParam);
   }
 
+  /**
+   * Build a SearchParam from Query.
+   *
+   * <pre>
+   * example usage:
+   * <code>
+   * SearchParam searchParam =
+   *     service.buildSearchParam(query)
+   *            .setParamsInJson("{\"fields\": [\"A\", \"B\"]}");
+   * </code>
+   * </pre>
+   */
   public SearchParam buildSearchParam(Query query) {
     return query.buildSearchParam(collectionName);
   }
+
+  /**
+   * Create a partition specified by <code>tag</code>
+   *
+   * @param tag partition tag
+   */
+  public void createPartition(String tag) {
+    client.createPartition(collectionName, tag);
+  }
+
+  /**
+   * Check whether the partition exists in this collection
+   *
+   * @param tag partition tag
+   * @return true if the partition exists, false otherwise.
+   */
+  public boolean hasPartition(String tag) {
+    return client.hasPartition(collectionName, tag);
+  }
+
+  /**
+   * List current partitions in the collection
+   *
+   * @return a list of partition names
+   */
+  public List<String> listPartitions() {
+    return client.listPartitions(collectionName);
+  }
+
+  /**
+   * Drop partition in the collection specified by <code>tag</code>
+   *
+   * @param tag partition tag
+   */
+  public void dropPartition(String tag) {
+    client.dropPartition(collectionName, tag);
+  }
+
+  /** Get current collection info */
+  public CollectionMapping getCollectionInfo() {
+    return client.getCollectionInfo(collectionName);
+  }
+
+  /** List collections. */
+  public List<String> listCollections() {
+    return client.listCollections();
+  }
+
+  /** Drop index by schema. */
+  public void dropIndex(Schema.VectorField vectorField) {
+    client.dropIndex(collectionName, vectorField.name);
+  }
+
+  /** Get collection stats. */
+  public String getCollectionStats() {
+    return client.getCollectionStats(collectionName);
+  }
+
+  /**
+   * Compacts the collection, erasing deleted data from disk and rebuild index in background (if the
+   * data size after compaction is still larger than indexFileSize). Data was only soft-deleted
+   * until you call compact.
+   */
+  public void compact() {
+    client.compact(CompactParam.create(collectionName));
+  }
+
+  /**
+   * Compacts the collection, erasing deleted data from disk and rebuild index in background (if the
+   * data size after compaction is still larger than indexFileSize). Data was only soft-deleted
+   * until you call compact.
+   */
+  public ListenableFuture<Void> compactAsync() {
+    return client.compactAsync(CompactParam.create(collectionName));
+  }
+
+  /**
+   * Compacts the collection, erasing deleted data from disk and rebuild index in background (if the
+   * data size after compaction is still larger than indexFileSize). Data was only soft-deleted
+   * until you call compact.
+   *
+   * @param threshold Defaults to 0.2. Segment will compact if and only if the percentage of
+   *     entities deleted exceeds the threshold.
+   */
+  public void compact(double threshold) {
+    client.compact(CompactParam.create(collectionName).setThreshold(threshold));
+  }
+
+  /**
+   * Compacts the collection, erasing deleted data from disk and rebuild index in background (if the
+   * data size after compaction is still larger than indexFileSize). Data was only soft-deleted
+   * until you call compact.
+   *
+   * @param threshold Defaults to 0.2. Segment will compact if and only if the percentage of
+   *     entities deleted exceeds the threshold.
+   */
+  public ListenableFuture<Void> compactAsync(double threshold) {
+    return client.compactAsync(CompactParam.create(collectionName).setThreshold(threshold));
+  }
 }

src/main/java/io/milvus/client/dsl/RangeQuery.java

@@ -5,6 +5,7 @@ import java.util.ArrayList;
 import java.util.List;
 import org.json.JSONObject;
 
+/** Range Query */
 public class RangeQuery<T> extends Query {
   private final Schema.Field<T> field;
   private final List<Expr> exprs = new ArrayList<>();
@@ -44,6 +45,13 @@ public class RangeQuery<T> extends Query {
     return json;
   }
 
+  /**
+   * Range query types.
+   * GT: greater than
+   * GTE: greater than or equal to
+   * LT: less than
+   * LTE: less than or equal to
+   */
   public enum Type {
     GT,
     GTE,

src/main/java/io/milvus/client/dsl/Schema.java

@@ -8,6 +8,7 @@ import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 
+/** An abstract class that allows you to predefine a schema to be used. */
 public abstract class Schema {
   private final Map<String, Field> fields = new LinkedHashMap<>();
 

src/main/java/io/milvus/client/dsl/TermQuery.java

@@ -5,6 +5,7 @@ import java.util.Collection;
 import org.json.JSONArray;
 import org.json.JSONObject;
 
+/** Term Query */
 public class TermQuery<T> extends Query {
   private final Schema.Field<T> field;
   private final Type type;
@@ -21,6 +22,10 @@ public class TermQuery<T> extends Query {
     return outer.put("term", new JSONObject().put(field.name, type.toJson(param)));
   }
 
+  /**
+   * Term query types.
+   * IN: field value should belong to one of the following values
+   */
   enum Type {
     IN {
       @Override

src/main/java/io/milvus/client/dsl/VectorQuery.java

@@ -11,6 +11,7 @@ import java.util.List;
 import java.util.stream.Collectors;
 import org.json.JSONObject;
 
+/** Vector query */
 public class VectorQuery<T> extends Query {
   private final Schema.VectorField<T> field;
   private final List<T> queries;
@@ -32,6 +33,7 @@ public class VectorQuery<T> extends Query {
     return this;
   }
 
+  /** topK */
   public VectorQuery<T> top(int topK) {
     this.topK = topK;
     return this;