Add more wrappers (#232)

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

examples/main/io/milvus/GeneralExample.java

@@ -459,7 +459,7 @@ public class GeneralExample {
         Random ran = new Random();
         for (int i = 0; i < 100; ++i) {
             R<MutationResult> result = example.insert(partitionName, row_count);
-            InsertResultWrapper wrapper = new InsertResultWrapper(result.getData());
+            MutationResultWrapper wrapper = new MutationResultWrapper(result.getData());
             List<Long> ids = wrapper.getLongIDs();
             deleteIds.add(ids.get(ran.nextInt(row_count)));
         }

src/main/java/io/milvus/Response/DescCollResponseWrapper.java

@@ -13,7 +13,7 @@ import java.util.ArrayList;
 import java.util.List;
 
 /**
- * Util class to wrap response of <code>DescribeCollection</code> interface.
+ * Util class to wrap response of <code>describeCollection</code> interface.
  */
 public class DescCollResponseWrapper {
     private final DescribeCollectionResponse response;

src/main/java/io/milvus/Response/DescIndexResponseWrapper.java

@@ -0,0 +1,98 @@
+package io.milvus.Response;
+
+import io.milvus.grpc.IndexDescription;
+import io.milvus.grpc.DescribeIndexResponse;
+
+import lombok.Getter;
+import lombok.NonNull;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Util class to wrap response of <code>describeIndex</code> interface.
+ */
+public class DescIndexResponseWrapper {
+    private final DescribeIndexResponse response;
+
+    public DescIndexResponseWrapper(@NonNull DescribeIndexResponse response) {
+        this.response = response;
+    }
+
+    /**
+     * Get index description of fields.
+     *
+     * @return <code>List<IndexDesc></code> index description of fields
+     */
+    public List<IndexDesc> GetIndexDescriptions() {
+        List<IndexDesc> results = new ArrayList<>();
+        List<IndexDescription> descriptions = response.getIndexDescriptionsList();
+        descriptions.forEach((desc)->{
+            IndexDesc res = new IndexDesc(desc.getFieldName(), desc.getIndexName(), desc.getIndexID());
+            desc.getParamsList().forEach((kv)-> res.AddParam(kv.getKey(), kv.getValue()));
+            results.add(res);
+        });
+
+        return results;
+    }
+
+    /**
+     * Get index description by field name.
+     * Return null if the field doesn't exist
+     *
+     * @return <code>IndexDesc</code> description of the index
+     */
+    public IndexDesc GetIndexDescByFieldName(@NonNull String name) {
+        for (int i = 0; i < response.getIndexDescriptionsCount(); ++i) {
+            IndexDescription desc = response.getIndexDescriptions(i);
+            if (name.compareTo(desc.getFieldName()) == 0) {
+                IndexDesc res = new IndexDesc(desc.getFieldName(), desc.getIndexName(), desc.getIndexID());
+                desc.getParamsList().forEach((kv)-> res.AddParam(kv.getKey(), kv.getValue()));
+                return res;
+            }
+        }
+
+        return null;
+    }
+
+    /**
+     * Internal-use class to wrap response of <code>describeIndex</code> interface.
+     */
+    @Getter
+    public static final class IndexDesc {
+        private final String fieldName;
+        private final String indexName;
+        private final long id;
+        private final Map<String, String> params = new HashMap<>();
+
+        public IndexDesc(@NonNull String fieldName, String indexName, long id) {
+            this.fieldName = fieldName;
+            this.indexName = indexName;
+            this.id = id;
+        }
+
+        public void AddParam(@NonNull String key, @NonNull String value) {
+            this.params.put(key, value);
+        }
+
+        @Override
+        public String toString() {
+            return "(fieldName: " + fieldName + " indexName: " + indexName + " id: " + id + " params: " +
+                    params.toString() + ")";
+        }
+    }
+
+    /**
+     * Construct a <code>String</code> by <code>DescIndexResponseWrapper</code> instance.
+     *
+     * @return <code>String</code>
+     */
+    @Override
+    public String toString() {
+        return "Index description{" +
+                GetIndexDescriptions().toString() +
+                '}';
+    }
+}

src/main/java/io/milvus/Response/GetPartStatResponseWrapper.java

@@ -0,0 +1,35 @@
+package io.milvus.Response;
+
+import io.milvus.grpc.GetPartitionStatisticsResponse;
+import io.milvus.grpc.KeyValuePair;
+import lombok.NonNull;
+
+import java.util.List;
+
+/**
+ * Utility class to wrap response of <code>getPartitionStatistics</code> interface.
+ */
+public class GetPartStatResponseWrapper {
+    private final GetPartitionStatisticsResponse stat;
+
+    public GetPartStatResponseWrapper(@NonNull GetPartitionStatisticsResponse stat) {
+        this.stat = stat;
+    }
+
+    /**
+     * Gets the row count of a field.
+     * Throw {@link NumberFormatException} if the row count is not a number.
+     *
+     * @return <code>int</code> dimension of the vector field
+     */
+    public long GetRowCount() throws NumberFormatException {
+        List<KeyValuePair> stats = stat.getStatsList();
+        for (KeyValuePair kv : stats) {
+            if (kv.getKey().compareTo("row_count") == 0) {
+                return Long.parseLong(kv.getValue());
+            }
+        }
+
+        return 0;
+    }
+}

src/main/java/io/milvus/Response/InsertResultWrapper.java → src/main/java/io/milvus/Response/MutationResultWrapper.java

@@ -8,12 +8,12 @@ import java.util.List;
 import lombok.NonNull;
 
 /**
- * Utility class to wrap response of <code>insert</code> interface.
+ * Utility class to wrap response of <code>insert/delete</code> interface.
  */
-public class InsertResultWrapper {
+public class MutationResultWrapper {
     private final MutationResult result;
 
-    public InsertResultWrapper(@NonNull MutationResult result) {
+    public MutationResultWrapper(@NonNull MutationResult result) {
         this.result = result;
     }
 
@@ -54,4 +54,13 @@ public class InsertResultWrapper {
             throw new ParamException("The primary key is not string type, please try getLongIDs()");
         }
     }
+
+    /**
+     * Gets the row count of the deleted entities. Currently this value is always equal to input row count
+     *
+     * @return <code>int</code> row count of the deleted entities
+     */
+    public long getDeleteCount() {
+        return result.getDeleteCnt();
+    }
 }

src/main/java/io/milvus/Response/ShowCollResponseWrapper.java

@@ -0,0 +1,105 @@
+package io.milvus.Response;
+
+import io.milvus.exception.IllegalResponseException;
+import io.milvus.grpc.ShowCollectionsResponse;
+
+import lombok.Getter;
+import lombok.NonNull;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * Util class to wrap response of <code>showCollections</code> interface.
+ */
+public class ShowCollResponseWrapper {
+    private final ShowCollectionsResponse response;
+
+    public ShowCollResponseWrapper(@NonNull ShowCollectionsResponse response) {
+        this.response = response;
+    }
+
+    /**
+     * Get information of the collections.
+     *
+     * @return <code>List<CollectionInfo></code> information array of the collections
+     */
+    public List<CollectionInfo> GetCollectionsInfo() throws IllegalResponseException {
+        if (response.getCollectionNamesCount() != response.getCollectionIdsCount()
+            || response.getCollectionNamesCount() != response.getCreatedUtcTimestampsCount()) {
+            throw new IllegalResponseException("Collection information count doesn't match");
+        }
+
+        List<CollectionInfo> results = new ArrayList<>();
+        for (int i = 0; i < response.getCollectionNamesCount(); ++i) {
+            CollectionInfo info = new CollectionInfo(response.getCollectionNames(i), response.getCollectionIds(i),
+                    response.getCreatedUtcTimestamps(i));
+            if (response.getInMemoryPercentagesCount() > i) {
+                info.SetInMemoryPercentage(response.getInMemoryPercentages(i));
+            }
+            results.add(info);
+        }
+
+        return results;
+    }
+
+    /**
+     * Get information of one collection by name.
+     *
+     * @return <code>CollectionInfo</code> information of the collection
+     */
+    public CollectionInfo GetCollectionInfo(@NonNull String name) {
+        for (int i = 0; i < response.getCollectionNamesCount(); ++i) {
+            if ( name.compareTo(response.getCollectionNames(i)) == 0) {
+                CollectionInfo info = new CollectionInfo(response.getCollectionNames(i), response.getCollectionIds(i),
+                        response.getCreatedUtcTimestamps(i));
+                if (response.getInMemoryPercentagesCount() > i) {
+                    info.SetInMemoryPercentage(response.getInMemoryPercentages(i));
+                }
+                return info;
+            }
+        }
+
+        return null;
+    }
+
+    /**
+     * Internal-use class to wrap response of <code>ShowCollections</code> interface.
+     */
+    @Getter
+    public static final class CollectionInfo {
+        private final String name;
+        private final long id;
+        private final long utcTimestamp;
+        private long inMemoryPercentage = 0;
+
+        public CollectionInfo(String name, long id, long utcTimestamp) {
+            this.name = name;
+            this.id = id;
+            this.utcTimestamp = utcTimestamp;
+        }
+
+        public void SetInMemoryPercentage(long inMemoryPercentage) {
+            this.inMemoryPercentage = inMemoryPercentage;
+        }
+
+        @Override
+        public String toString() {
+            return "(name: " + name + " id: " + id + " utcTimestamp: " + utcTimestamp + " inMemoryPercentage: "
+                    + inMemoryPercentage + ")";
+        }
+    }
+
+    /**
+     * Construct a <code>String</code> by <code>ShowCollResponseWrapper</code> instance.
+     *
+     * @return <code>String</code>
+     */
+    @Override
+    public String toString() {
+        return "Collections{" +
+                GetCollectionsInfo().toString() +
+                '}';
+    }
+}
+

src/main/java/io/milvus/Response/ShowPartResponseWrapper.java

@@ -0,0 +1,105 @@
+package io.milvus.Response;
+
+import io.milvus.exception.IllegalResponseException;
+import io.milvus.grpc.ShowPartitionsResponse;
+
+import lombok.Getter;
+import lombok.NonNull;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * Util class to wrap response of <code>showPartitions</code> interface.
+ */
+public class ShowPartResponseWrapper {
+    private final ShowPartitionsResponse response;
+
+    public ShowPartResponseWrapper(@NonNull ShowPartitionsResponse response) {
+        this.response = response;
+    }
+
+    /**
+     * Get information of the partitions.
+     *
+     * @return <code>List<PartitionInfo></code> information array of the partitions
+     */
+    public List<PartitionInfo> GetPartitionsInfo() throws IllegalResponseException {
+        if (response.getPartitionNamesCount() != response.getPartitionIDsCount()
+                || response.getPartitionNamesCount() != response.getCreatedUtcTimestampsCount()) {
+            throw new IllegalResponseException("Partition information count doesn't match");
+        }
+
+        List<PartitionInfo> results = new ArrayList<>();
+        for (int i = 0; i < response.getPartitionNamesCount(); ++i) {
+            PartitionInfo info = new PartitionInfo(response.getPartitionNames(i), response.getPartitionIDs(i),
+                    response.getCreatedUtcTimestamps(i));
+            if (response.getInMemoryPercentagesCount() > i) {
+                info.SetInMemoryPercentage(response.getInMemoryPercentages(i));
+            }
+            results.add(info);
+        }
+
+        return results;
+    }
+
+    /**
+     * Get information of one partition by name.
+     *
+     * @return <code>PartitionInfo</code> information of the partition
+     */
+    public PartitionInfo GetPartitionInfo(@NonNull String name) {
+        for (int i = 0; i < response.getPartitionNamesCount(); ++i) {
+            if ( name.compareTo(response.getPartitionNames(i)) == 0) {
+                PartitionInfo info = new PartitionInfo(response.getPartitionNames(i), response.getPartitionIDs(i),
+                        response.getCreatedUtcTimestamps(i));
+                if (response.getInMemoryPercentagesCount() > i) {
+                    info.SetInMemoryPercentage(response.getInMemoryPercentages(i));
+                }
+                return info;
+            }
+        }
+
+        return null;
+    }
+
+    /**
+     * Internal-use class to wrap response of <code>ShowPartitions</code> interface.
+     */
+    @Getter
+    public static final class PartitionInfo {
+        private final String name;
+        private final long id;
+        private final long utcTimestamp;
+        private long inMemoryPercentage = 0;
+
+        public PartitionInfo(String name, long id, long utcTimestamp) {
+            this.name = name;
+            this.id = id;
+            this.utcTimestamp = utcTimestamp;
+        }
+
+        public void SetInMemoryPercentage(long inMemoryPercentage) {
+            this.inMemoryPercentage = inMemoryPercentage;
+        }
+
+        @Override
+        public String toString() {
+            return "(name: " + name + " id: " + id + " utcTimestamp: " + utcTimestamp + " inMemoryPercentage: "
+                    + inMemoryPercentage + ")";
+        }
+    }
+
+    /**
+     * Construct a <code>String</code> by <code>ShowPartResponseWrapper</code> instance.
+     *
+     * @return <code>String</code>
+     */
+    @Override
+    public String toString() {
+        return "Partitions{" +
+                GetPartitionsInfo().toString() +
+                '}';
+    }
+}
+

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

@@ -922,6 +922,17 @@ public abstract class AbstractMilvusGrpcClient implements MilvusClient {
         logInfo(requestParam.toString());
 
         try {
+            // flush collection if client command to do it(some times user may want to know the newest row count)
+            if (requestParam.isFlushCollection()) {
+                R<FlushResponse> response = flush(FlushParam.newBuilder()
+                        .addCollectionName(requestParam.getCollectionName())
+                        .withSyncFlush(Boolean.TRUE)
+                        .build());
+                if (response.getStatus() != R.Status.Success.getCode()) {
+                    return R.failed(R.Status.valueOf(response.getStatus()), response.getMessage());
+                }
+            }
+
             GetPartitionStatisticsRequest getPartitionStatisticsRequest = GetPartitionStatisticsRequest.newBuilder()
                     .setCollectionName(requestParam.getCollectionName())
                     .setPartitionName(requestParam.getPartitionName())

src/main/java/io/milvus/param/partition/GetPartitionStatisticsParam.java

@@ -32,10 +32,12 @@ import lombok.NonNull;
 public class GetPartitionStatisticsParam {
     private final String collectionName;
     private final String partitionName;
+    private final boolean flushCollection;
 
     private GetPartitionStatisticsParam(@NonNull Builder builder) {
         this.collectionName = builder.collectionName;
         this.partitionName = builder.partitionName;
+        this.flushCollection = builder.flushCollection;
     }
 
     public static Builder newBuilder() {
@@ -49,6 +51,10 @@ public class GetPartitionStatisticsParam {
         private String collectionName;
         private String partitionName;
 
+        // if flushCollection is true, getCollectionStatistics() firstly call flush() and wait flush() finish
+        // Note: use default interval and timeout to wait flush()
+        private Boolean flushCollection = Boolean.TRUE;
+
         private Builder() {
         }
 
@@ -64,7 +70,7 @@ public class GetPartitionStatisticsParam {
         }
 
         /**
-         * Sets thep artition name. Partition name cannot be empty or null.
+         * Sets the partition name. Partition name cannot be empty or null.
          *
          * @param partitionName partition name
          * @return <code>Builder</code>
@@ -74,6 +80,17 @@ public class GetPartitionStatisticsParam {
             return this;
         }
 
+        /**
+         * Requires a flush action before retrieving partition statistics.
+         *
+         * @param flush <code>Boolean.TRUE</code> require a flush action
+         * @return <code>Builder</code>
+         */
+        public Builder withFlush(@NonNull Boolean flush) {
+            this.flushCollection = flush;
+            return this;
+        }
+
         /**
          * Verifies parameters and creates a new <code>GetPartitionStatisticsParam</code> instance.
          *
@@ -97,6 +114,7 @@ public class GetPartitionStatisticsParam {
         return "GetPartitionStatisticsParam{" +
                 "collectionName='" + collectionName + '\'' +
                 ", partitionName='" + partitionName + '\'' +
+                ", flush=" + flushCollection +
                 '}';
     }
 }

src/test/java/io/milvus/client/MilvusClientDockerTest.java

@@ -28,6 +28,9 @@ import io.milvus.param.dml.QueryParam;
 import io.milvus.param.dml.SearchParam;
 
 import io.milvus.param.index.CreateIndexParam;
+import io.milvus.param.index.DescribeIndexParam;
+import io.milvus.param.partition.GetPartitionStatisticsParam;
+import io.milvus.param.partition.ShowPartitionsParam;
 import org.apache.commons.text.RandomStringGenerator;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
@@ -266,8 +269,20 @@ public class MilvusClientDockerTest {
         GetCollStatResponseWrapper stat = new GetCollStatResponseWrapper(statR.getData());
         System.out.println("Collection row count: " + stat.GetRowCount());
 
+        // get partition statistics
+        R<GetPartitionStatisticsResponse> statPartR = client.getPartitionStatistics(GetPartitionStatisticsParam
+                .newBuilder()
+                .withCollectionName(randomCollectionName)
+                .withPartitionName("_default") // each collection has '_default' partition
+                .withFlush(true)
+                .build());
+        assertEquals(statPartR.getStatus().intValue(), R.Status.Success.getCode());
+
+        GetPartStatResponseWrapper statPart = new GetPartStatResponseWrapper(statPartR.getData());
+        System.out.println("Partition row count: " + statPart.GetRowCount());
+
         // create index
-        CreateIndexParam param = CreateIndexParam.newBuilder()
+        CreateIndexParam indexParam = CreateIndexParam.newBuilder()
                 .withCollectionName(randomCollectionName)
                 .withFieldName(field2Name)
                 .withIndexType(IndexType.IVF_FLAT)
@@ -278,15 +293,43 @@ public class MilvusClientDockerTest {
                 .withSyncWaitingTimeout(30L)
                 .build();
 
-        R<RpcStatus> createIndexR = client.createIndex(param);
+        R<RpcStatus> createIndexR = client.createIndex(indexParam);
         assertEquals(createIndexR.getStatus().intValue(), R.Status.Success.getCode());
 
+        // get index description
+        DescribeIndexParam descIndexParam = DescribeIndexParam.newBuilder()
+                .withCollectionName(randomCollectionName)
+                .withFieldName(field2Name)
+                .build();
+        R<DescribeIndexResponse> descIndexR = client.describeIndex(descIndexParam);
+        assertEquals(descIndexR.getStatus().intValue(), R.Status.Success.getCode());
+
+        DescIndexResponseWrapper indexDesc = new DescIndexResponseWrapper(descIndexR.getData());
+        System.out.println("Index description: " + indexDesc.toString());
+
         // load collection
         R<RpcStatus> loadR = client.loadCollection(LoadCollectionParam.newBuilder()
                 .withCollectionName(randomCollectionName)
                 .build());
         assertEquals(loadR.getStatus().intValue(), R.Status.Success.getCode());
 
+        // show collections
+        R<ShowCollectionsResponse> showR = client.showCollections(ShowCollectionsParam.newBuilder()
+                .addCollectionName(randomCollectionName)
+                .build());
+        assertEquals(showR.getStatus().intValue(), R.Status.Success.getCode());
+        ShowCollResponseWrapper info = new ShowCollResponseWrapper(showR.getData());
+        System.out.println("Collection info: " + info.toString());
+
+        // show partitions
+        R<ShowPartitionsResponse> showPartR = client.showPartitions(ShowPartitionsParam.newBuilder()
+                .withCollectionName(randomCollectionName)
+                .addPartitionName("_default") // each collection has a '_default' partition
+                .build());
+        assertEquals(showPartR.getStatus().intValue(), R.Status.Success.getCode());
+        ShowPartResponseWrapper infoPart = new ShowPartResponseWrapper(showPartR.getData());
+        System.out.println("Partition info: " + infoPart.toString());
+
         // query vectors to verify
         List<Long> queryIDs = new ArrayList<>();
         List<Boolean> compareGenders = new ArrayList<>();
@@ -441,7 +484,7 @@ public class MilvusClientDockerTest {
         R<MutationResult> insertR = client.insert(insertParam);
         assertEquals(insertR.getStatus().intValue(), R.Status.Success.getCode());
 //        System.out.println(insertR.getData());
-        InsertResultWrapper insertResultWrapper = new InsertResultWrapper(insertR.getData());
+        MutationResultWrapper insertResultWrapper = new MutationResultWrapper(insertR.getData());
         System.out.println(insertResultWrapper.getInsertCount() + " rows inserted");
         List<Long> ids = insertResultWrapper.getLongIDs();
 //        System.out.println("Auto-generated ids: " + ids);

src/test/java/io/milvus/client/MilvusServiceClientTest.java

@@ -1982,21 +1982,22 @@ class MilvusServiceClientTest {
         response = GetCollectionStatisticsResponse.newBuilder().build();
         wrapper = new GetCollStatResponseWrapper(response);
         assertEquals(0, wrapper.GetRowCount());
-
     }
 
     @Test
-    void InsertResultWrapper() {
+    void MutationResultWrapper() {
         List<Long> nID = Arrays.asList(1L, 2L, 3L);
         MutationResult results = MutationResult.newBuilder()
                 .setInsertCnt(nID.size())
+                .setDeleteCnt(nID.size())
                 .setIDs(IDs.newBuilder()
                         .setIntId(LongArray.newBuilder()
                                 .addAllData(nID)
                                 .build()))
                 .build();
-        InsertResultWrapper longWrapper = new InsertResultWrapper(results);
+        MutationResultWrapper longWrapper = new MutationResultWrapper(results);
         assertEquals(nID.size(), longWrapper.getInsertCount());
+        assertEquals(nID.size(), longWrapper.getDeleteCount());
         assertThrows(ParamException.class, longWrapper::getStringIDs);
 
         List<Long> longIDs = longWrapper.getLongIDs();
@@ -2013,7 +2014,7 @@ class MilvusServiceClientTest {
                                 .addAllData(sID)
                                 .build()))
                 .build();
-        InsertResultWrapper strWrapper = new InsertResultWrapper(results);
+        MutationResultWrapper strWrapper = new MutationResultWrapper(results);
         assertEquals(sID.size(), strWrapper.getInsertCount());
         assertThrows(ParamException.class, strWrapper::getLongIDs);