Support aggregation for search() (#1936) (#1938)

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

Author: yhmo

examples/src/main/java/io/milvus/v2/SearchAggregationExample.java

@@ -120,6 +120,7 @@ private SearchReq copy(SearchReq request) {
                 .functionScore(request.getFunctionScore())
                 .filterTemplateValues(request.getFilterTemplateValues())
                 .highlighter(request.getHighlighter())
+                .searchAggregation(request.getSearchAggregation())
                 .build();
     }
 

sdk-core/src/main/java/io/milvus/v2/client/MilvusClientV2Session.java

@@ -268,6 +268,7 @@ public SearchResp search(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStu
 
         SearchResp.SearchRespBuilder respBuilder = SearchResp.builder()
                 .searchResults(convertUtils.getEntities(response))
+                .aggregationBuckets(convertUtils.getAggregationBuckets(response))
                 .sessionTs(response.getSessionTs())
                 .recalls(response.getResults().getRecallsList());
         fillSearchRespFromExtraInfo(respBuilder, response.getStatus().getExtraInfoMap());

sdk-core/src/main/java/io/milvus/v2/service/vector/VectorService.java

@@ -22,6 +22,7 @@
 import io.milvus.v2.common.ConsistencyLevel;
 import io.milvus.v2.common.IndexParam;
 import io.milvus.v2.service.collection.request.CreateCollectionReq;
+import io.milvus.v2.service.vector.request.aggregation.SearchAggregation;
 import io.milvus.v2.service.vector.request.data.BaseVector;
 import io.milvus.v2.service.vector.request.highlighter.Highlighter;
 
@@ -74,6 +75,8 @@ public class SearchReq {
     // milvus v2.6.9 supports highlighter for search results
     private Highlighter highlighter;
 
+    private SearchAggregation searchAggregation;
+
     private SearchReq(SearchReqBuilder builder) {
         this.databaseName = builder.databaseName;
         this.collectionName = builder.collectionName;
@@ -102,6 +105,7 @@ private SearchReq(SearchReqBuilder builder) {
         this.filterTemplateValues = builder.filterTemplateValues;
         this.timezone = builder.timezone;
         this.highlighter = builder.highlighter;
+        this.searchAggregation = builder.searchAggregation;
     }
 
     // Getters and Setters
@@ -313,6 +317,14 @@ public Highlighter getHighlighter() {
         return highlighter;
     }
 
+    public SearchAggregation getSearchAggregation() {
+        return searchAggregation;
+    }
+
+    public void setSearchAggregation(SearchAggregation searchAggregation) {
+        this.searchAggregation = searchAggregation;
+    }
+
     @Override
     public String toString() {
         return "SearchReq{" +
@@ -340,6 +352,7 @@ public String toString() {
                 ", strictGroupSize=" + strictGroupSize +
                 ", ranker=" + ranker +
                 ", highlighter=" + (highlighter == null ? "null" : (highlighter.highlightType() + ":" + highlighter.getParams())) +
+                ", searchAggregation=" + searchAggregation +
                 ", functionScore=" + functionScore +
 //                ", filterTemplateValues=" + filterTemplateValues +
                 '}';
@@ -377,6 +390,7 @@ public static class SearchReqBuilder {
         private FunctionScore functionScore;
         private Map<String, Object> filterTemplateValues = new HashMap<>(); // default value
         private Highlighter highlighter;
+        private SearchAggregation searchAggregation;
 
         private SearchReqBuilder() {
         }
@@ -520,6 +534,11 @@ public SearchReqBuilder highlighter(Highlighter highlighter) {
             return this;
         }
 
+        public SearchReqBuilder searchAggregation(SearchAggregation searchAggregation) {
+            this.searchAggregation = searchAggregation;
+            return this;
+        }
+
         public SearchReq build() {
             return new SearchReq(this);
         }

sdk-core/src/main/java/io/milvus/v2/service/vector/request/SearchReq.java

@@ -0,0 +1,48 @@
+/*
+ * 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.v2.service.vector.request.aggregation;
+
+import io.milvus.v2.exception.ErrorCode;
+import io.milvus.v2.exception.MilvusClientException;
+
+public enum AggDirection {
+    ASC("asc"),
+    DESC("desc");
+
+    private final String value;
+
+    AggDirection(String value) {
+        this.value = value;
+    }
+
+    public String getValue() {
+        return value;
+    }
+
+    public static AggDirection fromValue(String value, String fieldName) {
+        for (AggDirection direction : values()) {
+            if (direction.value.equals(value)) {
+                return direction;
+            }
+        }
+        throw new MilvusClientException(ErrorCode.INVALID_PARAMS,
+                fieldName + " must be 'asc' or 'desc'.");
+    }
+}

sdk-core/src/main/java/io/milvus/v2/service/vector/request/aggregation/AggDirection.java

@@ -0,0 +1,51 @@
+/*
+ * 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.v2.service.vector.request.aggregation;
+
+import io.milvus.v2.exception.ErrorCode;
+import io.milvus.v2.exception.MilvusClientException;
+
+public enum MetricOps {
+    AVG("avg"),
+    SUM("sum"),
+    COUNT("count"),
+    MIN("min"),
+    MAX("max");
+
+    private final String value;
+
+    MetricOps(String value) {
+        this.value = value;
+    }
+
+    public String getValue() {
+        return value;
+    }
+
+    public static MetricOps fromValue(String value) {
+        for (MetricOps op : values()) {
+            if (op.value.equals(value)) {
+                return op;
+            }
+        }
+        throw new MilvusClientException(ErrorCode.INVALID_PARAMS,
+                "SearchAggregation metric op must be one of [avg, sum, count, min, max].");
+    }
+}

sdk-core/src/main/java/io/milvus/v2/service/vector/request/aggregation/MetricOps.java

@@ -0,0 +1,96 @@
+/*
+ * 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.v2.service.vector.request.aggregation;
+
+import io.milvus.grpc.MetricAggSpec;
+import io.milvus.v2.exception.ErrorCode;
+import io.milvus.v2.exception.MilvusClientException;
+
+public class MetricSpec {
+    private final MetricOps op;
+    private final String fieldName;
+
+    private MetricSpec(MetricSpecBuilder builder) {
+        if (builder.op == null) {
+            throw new MilvusClientException(ErrorCode.INVALID_PARAMS,
+                    "SearchAggregation metric op must not be null.");
+        }
+        if (builder.fieldName == null || builder.fieldName.isEmpty()) {
+            throw new MilvusClientException(ErrorCode.INVALID_PARAMS,
+                    "SearchAggregation metric fieldName must not be empty.");
+        }
+        if (builder.op != MetricOps.COUNT && "*".equals(builder.fieldName)) {
+            throw new MilvusClientException(ErrorCode.INVALID_PARAMS,
+                    "'*' is only valid for SearchAggregation count metrics.");
+        }
+        this.op = builder.op;
+        this.fieldName = builder.fieldName;
+    }
+
+    public static MetricSpecBuilder builder() {
+        return new MetricSpecBuilder();
+    }
+
+    public MetricOps getOp() {
+        return op;
+    }
+
+    public String getFieldName() {
+        return fieldName;
+    }
+
+    MetricAggSpec toProto() {
+        return MetricAggSpec.newBuilder()
+                .setOp(op.getValue())
+                .setFieldName(fieldName)
+                .build();
+    }
+
+    @Override
+    public String toString() {
+        return "MetricSpec{" +
+                "op=" + op +
+                ", fieldName='" + fieldName + '\'' +
+                '}';
+    }
+
+    public static class MetricSpecBuilder {
+        private MetricOps op;
+        private String fieldName;
+
+        private MetricSpecBuilder() {
+        }
+
+        public MetricSpecBuilder op(MetricOps op) {
+            this.op = op;
+            return this;
+        }
+
+        public MetricSpecBuilder fieldName(String fieldName) {
+            this.fieldName = fieldName;
+            return this;
+        }
+
+        public MetricSpec build() {
+            return new MetricSpec(this);
+        }
+    }
+
+}

sdk-core/src/main/java/io/milvus/v2/service/vector/request/aggregation/MetricSpec.java

@@ -0,0 +1,107 @@
+/*
+ * 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.v2.service.vector.request.aggregation;
+
+import io.milvus.v2.exception.ErrorCode;
+import io.milvus.v2.exception.MilvusClientException;
+
+public class OrderSpec {
+    private final String key;
+    private final AggDirection direction;
+    private final Boolean nullFirst;
+
+    private OrderSpec(OrderSpecBuilder builder) {
+        if (builder.key == null || builder.key.isEmpty()) {
+            throw new MilvusClientException(ErrorCode.INVALID_PARAMS,
+                    "SearchAggregation.order key must not be empty.");
+        }
+        if (builder.direction == null) {
+            throw new MilvusClientException(ErrorCode.INVALID_PARAMS,
+                    "SearchAggregation.order direction must not be null.");
+        }
+        this.key = builder.key;
+        this.direction = builder.direction;
+        this.nullFirst = builder.nullFirst;
+    }
+
+    public static OrderSpecBuilder builder() {
+        return new OrderSpecBuilder();
+    }
+
+    public String getKey() {
+        return key;
+    }
+
+    public AggDirection getDirection() {
+        return direction;
+    }
+
+    public Boolean getNullFirst() {
+        return nullFirst;
+    }
+
+    io.milvus.grpc.OrderSpec toProto() {
+        io.milvus.grpc.OrderSpec.Builder builder = io.milvus.grpc.OrderSpec.newBuilder()
+                .setKey(key)
+                .setDirection(direction.getValue());
+        if (nullFirst != null) {
+            builder.setNullFirst(nullFirst);
+        }
+        return builder.build();
+    }
+
+    @Override
+    public String toString() {
+        return "OrderSpec{" +
+                "key='" + key + '\'' +
+                ", direction=" + direction +
+                ", nullFirst=" + nullFirst +
+                '}';
+    }
+
+    public static class OrderSpecBuilder {
+        private String key;
+        private AggDirection direction;
+        private Boolean nullFirst;
+
+        private OrderSpecBuilder() {
+        }
+
+        public OrderSpecBuilder key(String key) {
+            this.key = key;
+            return this;
+        }
+
+        public OrderSpecBuilder direction(AggDirection direction) {
+            this.direction = direction;
+            return this;
+        }
+
+        public OrderSpecBuilder nullFirst(Boolean nullFirst) {
+            this.nullFirst = nullFirst;
+            return this;
+        }
+
+        public OrderSpec build() {
+            return new OrderSpec(this);
+        }
+    }
+
+}