Added size attribute to MultiTermsAggregation

Signed-off-by: Fabian Kopatschek <[email protected]>
opensearch-project · Sep 22, 2023 · 4eb9510 · 4eb9510
1 parent 66608d3
commit 4eb9510
Show file tree

Hide file tree

Showing 3 changed files with 104 additions and 2 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -7,6 +7,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 - Document HTTP/2 support ([#330](https://github.com/opensearch-project/opensearch-java/pull/330))
 - Added support for "cjk" analyzer ([#604](https://github.com/opensearch-project/opensearch-java/pull/604))
 - Added support for wrapper queries ([#630](https://github.com/opensearch-project/opensearch-java/pull/630))
+- Added size attribute to MultiTermsAggregation ([#627](https://github.com/opensearch-project/opensearch-java/pull/627))
 
 ### Dependencies
 - Bumps `org.ajoberstar.grgit:grgit-gradle` from 5.0.0 to 5.2.0

diff --git a/...main/java/org/opensearch/client/opensearch/_types/aggregations/MultiTermsAggregation.java b/...main/java/org/opensearch/client/opensearch/_types/aggregations/MultiTermsAggregation.java
@@ -43,19 +43,24 @@
 import java.util.List;
 import java.util.function.Function;
 
+import javax.annotation.Nullable;
+
 // typedef: _types.aggregations.MultiTermsAggregation
 
 @JsonpDeserializable
 public class MultiTermsAggregation extends BucketAggregationBase implements AggregationVariant {
 	private final List<MultiTermLookup> terms;
 
+	@Nullable
+	private final Integer size;
+
 	// ---------------------------------------------------------------------------------------------
 
 	private MultiTermsAggregation(Builder builder) {
 		super(builder);
 
 		this.terms = ApiTypeHelper.unmodifiableRequired(builder.terms, this, "terms");
-
+		this.size = builder.size;
 	}
 
 	public static MultiTermsAggregation of(Function<Builder, ObjectBuilder<MultiTermsAggregation>> fn) {
@@ -77,9 +82,22 @@ public final List<MultiTermLookup> terms() {
 		return this.terms;
 	}
 
+	/**
+	 * API name: {@code size}
+	 */
+	@Nullable
+	public final Integer size() {
+		return this.size;
+	}
+
 	protected void serializeInternal(JsonGenerator generator, JsonpMapper mapper) {
 
 		super.serializeInternal(generator, mapper);
+		if (this.size != null) {
+			generator.writeKey("size");
+			generator.write(this.size);
+
+		}
 		if (ApiTypeHelper.isDefined(this.terms)) {
 			generator.writeKey("terms");
 			generator.writeStartArray();
@@ -104,6 +122,9 @@ public static class Builder extends BucketAggregationBase.AbstractBuilder<Builde
 				ObjectBuilder<MultiTermsAggregation> {
 		private List<MultiTermLookup> terms;
 
+		@Nullable
+		private Integer size;
+
 		/**
 		 * Required - API name: {@code terms}
 		 * <p>
@@ -133,6 +154,14 @@ public final Builder terms(Function<MultiTermLookup.Builder, ObjectBuilder<Multi
 			return terms(fn.apply(new MultiTermLookup.Builder()).build());
 		}
 
+		/**
+		 * API name: {@code size}
+		 */
+		public final Builder size(@Nullable Integer value) {
+			this.size = value;
+			return this;
+		}
+
 		@Override
 		protected Builder self() {
 			return this;
@@ -161,8 +190,8 @@ public MultiTermsAggregation build() {
 
 	protected static void setupMultiTermsAggregationDeserializer(ObjectDeserializer<MultiTermsAggregation.Builder> op) {
 		BucketAggregationBase.setupBucketAggregationBaseDeserializer(op);
+		op.add(Builder::size, JsonpDeserializer.integerDeserializer(), "size");
 		op.add(Builder::terms, JsonpDeserializer.arrayDeserializer(MultiTermLookup._DESERIALIZER), "terms");
-
 	}
 
 }
diff --git a/...rc/test/java/org/opensearch/client/opensearch/integTest/AbstractAggregationRequestIT.java b/...rc/test/java/org/opensearch/client/opensearch/integTest/AbstractAggregationRequestIT.java
@@ -15,6 +15,8 @@
 import org.opensearch.client.opensearch._types.aggregations.DateRangeAggregation;
 import org.opensearch.client.opensearch._types.aggregations.DateRangeExpression;
 import org.opensearch.client.opensearch._types.aggregations.FieldDateMath;
+import org.opensearch.client.opensearch._types.aggregations.MultiTermLookup;
+import org.opensearch.client.opensearch._types.aggregations.MultiTermsAggregation;
 import org.opensearch.client.opensearch._types.aggregations.RangeAggregation;
 import org.opensearch.client.opensearch._types.mapping.Property;
 import org.opensearch.client.opensearch.core.SearchResponse;
@@ -63,6 +65,54 @@ public void testDateRangeAggregation() throws Exception {
 		assertEquals(2, buckets.get(2).docCount());
 	}
 
+	@Test
+	public void testMultiTermsAggregation() throws Exception {
+		var index = "test-multiterms-aggregation-no-size";
+		createMultiTermsDocuments(index);
+		var searchResponse = sendAggregateRequest(index, "multiterms", getMultiTermsAggregation(null));
+		var multitermsAggregations = searchResponse.aggregations().get("multiterms");
+		var buckets = multitermsAggregations._get()
+		._toAggregate()
+		.multiTerms()
+		.buckets()
+		.array();
+
+		assertEquals(3, buckets.size());
+		assertEquals(1, buckets.get(0).docCount());
+	}
+
+	@Test
+	public void testMultiTermsAggregationWithSizeFewerThenBucketsSize() throws Exception {
+		var index = "test-multiterms-aggregation-fewer-size";
+		createMultiTermsDocuments(index);
+		var searchResponse = sendAggregateRequest(index, "multiterms", getMultiTermsAggregation(1));
+		var multitermsAggregations = searchResponse.aggregations().get("multiterms");
+		var buckets = multitermsAggregations._get()
+				._toAggregate()
+				.multiTerms()
+				.buckets()
+				.array();
+
+		assertEquals(1, buckets.size());
+		assertEquals(1, buckets.get(0).docCount());
+	}
+
+	@Test
+	public void testMultiTermsAggregationWithSizeBiggerThenBucketsSize() throws Exception {
+		var index = "test-multiterms-aggregation-bigger-size";
+		createMultiTermsDocuments(index);
+		var searchResponse = sendAggregateRequest(index, "multiterms", getMultiTermsAggregation(50));
+		var multitermsAggregations = searchResponse.aggregations().get("multiterms");
+		var buckets = multitermsAggregations._get()
+				._toAggregate()
+				.multiTerms()
+				.buckets()
+				.array();
+
+		assertEquals(3, buckets.size());
+		assertEquals(1, buckets.get(0).docCount());
+	}
+
 	private Aggregation getExpiryDateRangeAggregation() {
 		DateRangeAggregation expiryDateRangeAggregation = new DateRangeAggregation.Builder()
 				.field("expDate")
@@ -79,6 +129,21 @@ private Aggregation getCostValueRangeAggregation() {
 		return new Aggregation.Builder().range(costValueRangeAggregation).build();
 	}
 
+	private Aggregation getMultiTermsAggregation(Integer size) {
+		MultiTermsAggregation.Builder multiTermsAggregationBuilder = new MultiTermsAggregation.Builder()
+				.terms(List.of(
+					MultiTermLookup.of(multiTermLookup -> multiTermLookup.field("cost")),
+					MultiTermLookup.of(multiTermLookup -> multiTermLookup.field("expDate"))
+				));
+
+		if(size != null) {
+			multiTermsAggregationBuilder = multiTermsAggregationBuilder.size(size);
+		}
+
+		MultiTermsAggregation multiTermsAggregation = multiTermsAggregationBuilder.build();
+		return new Aggregation.Builder().multiTerms(multiTermsAggregation).build();
+	}
+
 	private SearchResponse<Void> sendAggregateRequest(String index, String key, Aggregation value) throws IOException {
 		return javaClient().search(
 				request -> request.index(index)
@@ -131,6 +196,13 @@ private void createDateRangeDocuments(String index) throws IOException {
 		javaClient().create(_1 -> _1.index(index).id("6").document(createProduct("oil", 50, 6)).refresh(Refresh.True));
 	}
 
+	private void createMultiTermsDocuments(String index) throws IOException {
+		createIndex(index);
+		javaClient().create(_1 -> _1.index(index).id("1").document(createProduct("appleA", 2, 1)).refresh(Refresh.True));
+		javaClient().create(_1 -> _1.index(index).id("2").document(createProduct("appleB", 2, 2)).refresh(Refresh.True));
+		javaClient().create(_1 -> _1.index(index).id("3").document(createProduct("appleC", 2, 3)).refresh(Refresh.True));
+	}
+
 	private void createIndex(String index) throws IOException {
 		Property nameValueProp = new Property.Builder()
 				.text(v -> v)