diff --git a/src/it/java/io/weaviate/integration/AggregationITest.java b/src/it/java/io/weaviate/integration/AggregationITest.java index 9df1bc287..071607d66 100644 --- a/src/it/java/io/weaviate/integration/AggregationITest.java +++ b/src/it/java/io/weaviate/integration/AggregationITest.java @@ -80,7 +80,7 @@ public void testOverAll_groupBy_category() { Aggregation.integer("price", calculate -> calculate.min().max().count())) .includeTotalCount(true), - new GroupBy("category")); + GroupBy.property("category")); Assertions.assertThat(result) .extracting(AggregateResponseGrouped::groups) @@ -139,7 +139,7 @@ public void testNearVector_groupBy_category() { calculate -> calculate.min().max().median())) .objectLimit(9) .includeTotalCount(true), - new GroupBy("category")); + GroupBy.property("category")); Assertions.assertThat(result) .extracting(AggregateResponseGrouped::groups) diff --git a/src/main/java/io/weaviate/client6/v1/api/collections/aggregate/GroupBy.java b/src/main/java/io/weaviate/client6/v1/api/collections/aggregate/GroupBy.java index f60898a21..e84b74a8f 100644 --- a/src/main/java/io/weaviate/client6/v1/api/collections/aggregate/GroupBy.java +++ b/src/main/java/io/weaviate/client6/v1/api/collections/aggregate/GroupBy.java @@ -1,13 +1,48 @@ package io.weaviate.client6.v1.api.collections.aggregate; +import java.util.function.Function; + +import io.weaviate.client6.v1.internal.ObjectBuilder; import io.weaviate.client6.v1.internal.grpc.protocol.WeaviateProtoAggregate; -public record GroupBy(String property) { - public static final GroupBy of(String property) { - return new GroupBy(property); +public record GroupBy(String property, Integer limit) { + public static final GroupBy property(String property) { + return property(property, ObjectBuilder.identity()); + } + + public static final GroupBy property(String property, Function> fn) { + return fn.apply(new Builder(property)).build(); + } + + public GroupBy(Builder builder) { + this(builder.property, builder.limit); + } + + public static class Builder implements ObjectBuilder { + private final String property; + + public Builder(String property) { + this.property = property; + } + + private Integer limit; + + public final Builder limit(int limit) { + this.limit = limit; + return this; + } + + @Override + public GroupBy build() { + return new GroupBy(this); + } } void appendTo(WeaviateProtoAggregate.AggregateRequest.Builder req, String collection) { + if (limit != null) { + req.setLimit(limit); + } + req.setGroupBy(WeaviateProtoAggregate.AggregateRequest.GroupBy.newBuilder() .setCollection(collection) .setProperty(property));