Cardinality/Missing Aggs

gradle/libs.versions.toml

@@ -22,7 +22,7 @@ kotlin            = "1.7.10"
 antlr4            = "4.13.1"
 guava             = "33.2.1-jre"
 gson              = "2.13.2"
-opensearchprotobufs = "0.23.0"
+opensearchprotobufs = "0.234.0-SNAPSHOT"
 protobuf          = "3.25.8"
 jakarta_annotation = "1.3.5"
 google_http_client = "1.44.1"

modules/transport-grpc/build.gradle

@@ -6,6 +6,10 @@
  * compatible open source license.
  */
 
+repositories {
+  mavenLocal()
+}
+
 apply plugin: 'opensearch.testclusters'
 apply plugin: 'opensearch.internal-cluster-test'
 

modules/transport-grpc/spi/build.gradle

@@ -6,6 +6,10 @@
  * compatible open source license.
  */
 
+repositories {
+  mavenLocal()
+}
+
 apply plugin: 'opensearch.build'
 apply plugin: 'opensearch.publish'
 

modules/transport-grpc/src/internalClusterTest/java/org/opensearch/transport/grpc/DocumentServiceIT.java

@@ -44,7 +44,7 @@ public void testDocumentServiceBulk() throws Exception {
                 .setObject(com.google.protobuf.ByteString.copyFromUtf8(DEFAULT_DOCUMENT_SOURCE))
                 .build();
 
-            BulkRequest bulkRequest = BulkRequest.newBuilder().addBulkRequestBody(requestBody).build();
+            BulkRequest bulkRequest = BulkRequest.newBuilder().addRequestBody(requestBody).build();
 
             // Execute the bulk request
             BulkResponse bulkResponse = documentStub.bulk(bulkRequest);

modules/transport-grpc/src/internalClusterTest/java/org/opensearch/transport/grpc/SearchServiceIT.java

@@ -8,6 +8,11 @@
 
 package org.opensearch.transport.grpc;
 
+import org.opensearch.common.settings.Settings;
+import org.opensearch.common.xcontent.XContentType;
+import org.opensearch.protobufs.AggregationContainer;
+import org.opensearch.protobufs.CardinalityAggregation;
+import org.opensearch.protobufs.MissingAggregation;
 import org.opensearch.protobufs.SearchRequest;
 import org.opensearch.protobufs.SearchRequestBody;
 import org.opensearch.protobufs.SearchResponse;
@@ -16,6 +21,8 @@
 
 import io.grpc.ManagedChannel;
 
+import java.util.List;
+
 /**
  * Integration tests for the SearchService gRPC service.
  */
@@ -43,7 +50,7 @@ public void testSearchServiceSearch() throws Exception {
 
             SearchRequest searchRequest = SearchRequest.newBuilder()
                 .addIndex(indexName)
-                .setSearchRequestBody(requestBody)
+                .setRequestBody(requestBody)
                 .setQ("field1:value1")
                 .build();
 
@@ -57,4 +64,175 @@ public void testSearchServiceSearch() throws Exception {
             assertEquals("Hit should have correct ID", "1", searchResponse.getHits().getHits(0).getXId());
         }
     }
+
+    public void testBasicMissingAgg() throws Exception {
+        String indexName = "test-index";
+        createTestIndex(indexName);
+
+        List<String> testDocs = List.of(
+            "{\"species\":\"turtle\",\"age\":142}",
+            "{\"species\":\"cat\",\"age\":12}",
+            "{\"species\":\"dog\",\"age\":5}",
+            "{\"species\":\"dog\"}"
+        );
+        for (String testDoc : testDocs) {
+            client().prepareIndex(indexName).setSource(testDoc, XContentType.JSON).get();
+        }
+        client().admin().indices().prepareRefresh(indexName).get();
+
+        MissingAggregation.Builder missingAgg = MissingAggregation.newBuilder()
+            .setField("age");
+
+        AggregationContainer aggCont = AggregationContainer.newBuilder()
+            .setMissing(missingAgg)
+            .build();
+
+        SearchRequestBody requestBody = SearchRequestBody.newBuilder()
+            .putAggregations("missing_age", aggCont)
+            .build();
+
+        SearchRequest searchRequest = SearchRequest.newBuilder()
+            .addIndex(indexName)
+            .setRequestBody(requestBody)
+            .build();
+
+        SearchResponse searchResponse;
+        try (NettyGrpcClient client = createGrpcClient()) {
+            ManagedChannel channel = client.getChannel();
+            SearchServiceGrpc.SearchServiceBlockingStub searchStub = SearchServiceGrpc.newBlockingStub(channel);
+            searchResponse = searchStub.search(searchRequest);
+        }
+
+        assertNotNull("Search response should not be null", searchResponse);
+        assertEquals("Search response should return all documents as hits", 4, searchResponse.getHits().getHitsCount());
+        assertEquals("Missing doc count should be 1", 1, searchResponse.getAggregationsMap().get("missing_age").getMissing().getDocCount());
+    }
+
+    public void testBasicCardinalityAgg() throws Exception {
+        String indexName = "new-test-index";
+        String mapping = """
+            {
+              "properties": {
+                "species": {
+                  "type": "keyword"
+                },
+                "age": {
+                  "type": "integer"
+                }
+              }
+            }""";
+        createIndex(indexName, Settings.EMPTY, mapping);
+        ensureGreen(indexName);
+
+        List<String> testDocs = List.of(
+            "{\"species\":\"turtle\",\"age\":142}",
+            "{\"species\":\"cat\",\"age\":12}",
+            "{\"species\":\"dog\",\"age\":5}",
+            "{\"species\":\"dog\"}"
+        );
+        for (String testDoc : testDocs) {
+            client().prepareIndex(indexName).setSource(testDoc, XContentType.JSON).get();
+        }
+        client().admin().indices().prepareRefresh(indexName).get();
+
+        CardinalityAggregation.Builder cardinalityAgg = CardinalityAggregation.newBuilder()
+            .setField("species");
+
+        AggregationContainer aggCont = AggregationContainer.newBuilder()
+            .setCardinality(cardinalityAgg)
+            .build();
+
+        SearchRequestBody requestBody = SearchRequestBody.newBuilder()
+            .putAggregations("species_cardinality", aggCont)
+            .build();
+
+        SearchRequest searchRequest = SearchRequest.newBuilder()
+            .addIndex(indexName)
+            .setRequestBody(requestBody)
+            .build();
+
+        SearchResponse searchResponse;
+        try (NettyGrpcClient client = createGrpcClient()) {
+            ManagedChannel channel = client.getChannel();
+            SearchServiceGrpc.SearchServiceBlockingStub searchStub = SearchServiceGrpc.newBlockingStub(channel);
+            searchResponse = searchStub.search(searchRequest);
+        }
+
+        assertNotNull("Search response should not be null", searchResponse);
+        assertEquals("Search response should return all documents as hits", 4, searchResponse.getHits().getHitsCount());
+        assertEquals("Cardinality of species field should be 3", 3, searchResponse.getAggregationsMap().get("species_cardinality").getCardinality().getValue());
+    }
+
+    public void testMultiAggRequest() throws Exception {
+        String indexName = "new-test-index";
+        String mapping = """
+            {
+              "properties": {
+                "http-method": {
+                  "type": "keyword"
+                },
+                "response-code": {
+                  "type": "keyword"
+                },
+                "uri": {
+                  "type": "keyword"
+                }
+              }
+            }""";
+        createIndex(indexName, Settings.EMPTY, mapping);
+        ensureGreen(indexName);
+
+        List<String> testDocs = List.of(
+            "{\"http-method\":\"del\",\"response-code\":\"200\",\"uri\":142}",
+            "{\"http-method\":\"put\",\"response-code\":\"200\",\"uri\":142}",
+            "{\"http-method\":\"get\",\"response-code\":\"200\",\"uri\":142}",
+            "{\"http-method\":\"get\",\"response-code\":\"401\"}"
+        );
+        for (String testDoc : testDocs) {
+            client().prepareIndex(indexName).setSource(testDoc, XContentType.JSON).get();
+        }
+        client().admin().indices().prepareRefresh(indexName).get();
+
+        CardinalityAggregation.Builder methodCardAgg = CardinalityAggregation.newBuilder()
+            .setField("http-method");
+        AggregationContainer methodCardAggCont = AggregationContainer.newBuilder()
+            .setCardinality(methodCardAgg)
+            .build();
+
+        CardinalityAggregation.Builder respCardAgg = CardinalityAggregation.newBuilder()
+            .setField("response-code");
+        AggregationContainer respCardAggCont = AggregationContainer.newBuilder()
+            .setCardinality(respCardAgg)
+            .build();
+
+        MissingAggregation.Builder uriMissingAgg = MissingAggregation.newBuilder()
+            .setField("uri");
+        AggregationContainer uriMissingAggCont = AggregationContainer.newBuilder()
+            .setMissing(uriMissingAgg)
+            .build();
+
+        SearchRequestBody requestBody = SearchRequestBody.newBuilder()
+            .putAggregations("http-method-card", methodCardAggCont)
+            .putAggregations("http-response-card", respCardAggCont)
+            .putAggregations("no-uri", uriMissingAggCont)
+            .build();
+
+        SearchRequest searchRequest = SearchRequest.newBuilder()
+            .addIndex(indexName)
+            .setRequestBody(requestBody)
+            .build();
+
+        SearchResponse searchResponse;
+        try (NettyGrpcClient client = createGrpcClient()) {
+            ManagedChannel channel = client.getChannel();
+            SearchServiceGrpc.SearchServiceBlockingStub searchStub = SearchServiceGrpc.newBlockingStub(channel);
+            searchResponse = searchStub.search(searchRequest);
+        }
+
+        assertNotNull("Search response should not be null", searchResponse);
+        assertEquals("Search response should return all documents as hits", 4, searchResponse.getHits().getHitsCount());
+        assertEquals("Http method cardinality should be 3", 3, searchResponse.getAggregationsMap().get("http-method-card").getCardinality().getValue());
+        assertEquals("Http response code cardinality should be 2", 2, searchResponse.getAggregationsMap().get("http-response-card").getCardinality().getValue());
+        assertEquals("Documents missing uri should be 1", 1, searchResponse.getAggregationsMap().get("no-uri").getMissing().getDocCount());
+    }
 }

modules/transport-grpc/src/main/java/org/opensearch/transport/grpc/proto/request/document/bulk/BulkRequestParserProtoUtils.java

@@ -116,7 +116,7 @@ public static DocWriteRequest<?>[] getDocWriteRequests(
         String defaultPipeline,
         Boolean defaultRequireAlias
     ) {
-        List<BulkRequestBody> bulkRequestBodyList = request.getBulkRequestBodyList();
+        List<BulkRequestBody> bulkRequestBodyList = request.getRequestBodyList();
         DocWriteRequest<?>[] docWriteRequests = new DocWriteRequest<?>[bulkRequestBodyList.size()];
 
         // Process each operation in the request body

modules/transport-grpc/src/main/java/org/opensearch/transport/grpc/proto/request/search/SearchRequestProtoUtils.java

@@ -20,6 +20,7 @@
 import org.opensearch.rest.RestRequest;
 import org.opensearch.rest.action.search.RestSearchAction;
 import org.opensearch.search.Scroll;
+import org.opensearch.search.aggregations.AggregatorFactories;
 import org.opensearch.search.builder.SearchSourceBuilder;
 import org.opensearch.search.fetch.subphase.FetchSourceContext;
 import org.opensearch.search.internal.SearchContext;
@@ -113,7 +114,7 @@ protected static void parseSearchRequest(
         }
         searchRequest.indices(indexArr);
 
-        SearchSourceBuilderProtoUtils.parseProto(searchRequest.source(), request.getSearchRequestBody(), queryUtils);
+        SearchSourceBuilderProtoUtils.parseProto(searchRequest.source(), request.getRequestBody(), queryUtils);
 
         final int batchedReduceSize = request.hasBatchedReduceSize()
             ? request.getBatchedReduceSize()
@@ -287,7 +288,7 @@ private static void preparePointInTime(
      */
     protected static void checkProtoTotalHits(SearchRequest protoRequest, org.opensearch.action.search.SearchRequest searchRequest) {
 
-        boolean totalHitsAsInt = protoRequest.hasTotalHitsAsInt() ? protoRequest.getTotalHitsAsInt() : false;
+        boolean totalHitsAsInt = protoRequest.hasRestTotalHitsAsInt() ? protoRequest.getRestTotalHitsAsInt() : false;
         if (totalHitsAsInt == false) {
             return;
         }

modules/transport-grpc/src/main/java/org/opensearch/transport/grpc/proto/request/search/SearchSourceBuilderProtoUtils.java

@@ -9,19 +9,23 @@
 
 import org.opensearch.common.unit.TimeValue;
 import org.opensearch.core.xcontent.XContentParser;
+import org.opensearch.protobufs.AggregationContainer;
 import org.opensearch.protobufs.DerivedField;
 import org.opensearch.protobufs.FieldAndFormat;
 import org.opensearch.protobufs.Rescore;
 import org.opensearch.protobufs.ScriptField;
 import org.opensearch.protobufs.SearchRequestBody;
 import org.opensearch.protobufs.TrackHits;
+import org.opensearch.search.aggregations.Aggregation;
+import org.opensearch.search.aggregations.AggregationBuilder;
 import org.opensearch.search.builder.SearchSourceBuilder;
 import org.opensearch.search.sort.SortBuilder;
 import org.opensearch.transport.grpc.proto.request.common.FetchSourceContextProtoUtils;
 import org.opensearch.transport.grpc.proto.request.common.ScriptProtoUtils;
 import org.opensearch.transport.grpc.proto.request.search.query.AbstractQueryBuilderProtoUtils;
 import org.opensearch.transport.grpc.proto.request.search.sort.SortBuilderProtoUtils;
 import org.opensearch.transport.grpc.proto.request.search.suggest.SuggestBuilderProtoUtils;
+import org.opensearch.transport.grpc.proto.request.search.aggs.AggregationContainerBuilderProtoUtils;
 
 import java.io.IOException;
 import java.util.Map;
@@ -64,6 +68,18 @@ public static void parseProto(
         if (protoRequest.hasPostFilter()) {
             searchSourceBuilder.postFilter(queryUtils.parseInnerQueryBuilderProto(protoRequest.getPostFilter()));
         }
+
+        // Handle aggregations
+        Map<String, AggregationContainer> aggregationMap = protoRequest.getAggregationsMap();
+        for (Map.Entry<String, AggregationContainer> aggEntry : aggregationMap.entrySet()) {
+            // Convert aggregation if able
+            AggregationBuilder aggregationBuilder = AggregationContainerBuilderProtoUtils.fromProto(aggEntry);
+            // Ignore unrecognized or explicitly unspecified aggregations which return null.
+            if (aggregationBuilder != null) {
+                // Add aggregation entry to request source
+                searchSourceBuilder.aggregation(aggregationBuilder);
+            }
+        }
     }
 
     /**
@@ -142,11 +158,11 @@ private static void parseNonQueryFields(SearchSourceBuilder searchSourceBuilder,
                 searchSourceBuilder.scriptField(name, scriptField.script(), scriptField.ignoreFailure());
             }
         }
-        if (protoRequest.getIndicesBoostCount() > 0) {
-            for (Map.Entry<String, Float> entry : protoRequest.getIndicesBoostMap().entrySet()) {
-                searchSourceBuilder.indexBoost(entry.getKey(), entry.getValue());
-            }
-        }
+//        if (protoRequest.getIndicesBoostCount() > 0) {
+//            for (Map.Entry<String, Float> entry : protoRequest.getIndicesBoostList().entrySet()) {
+//                searchSourceBuilder.indexBoost(entry.getKey(), entry.getValue());
+//            }
+//        }
 
         // TODO support aggregations
         /*

modules/transport-grpc/src/main/java/org/opensearch/transport/grpc/proto/request/search/aggs/AggregationContainerBuilderProtoUtils.java

@@ -0,0 +1,50 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.transport.grpc.proto.request.search.aggs;
+
+import org.opensearch.protobufs.AggregationContainer;
+import org.opensearch.search.aggregations.AggregationBuilder;
+
+import java.io.IOException;
+import java.util.Map;
+
+/**
+ * Handle an aggregation container which wraps an aggregation type.
+ */
+public class AggregationContainerBuilderProtoUtils {
+
+    /**
+     * Private no-op.
+     */
+    private AggregationContainerBuilderProtoUtils() {
+        // Utility class, no instances
+    }
+
+    /**
+     * Maps aggregation types wrapped by AggregationContainer to the appropriate fromProto conversion.
+     * @throws IOException if there's an error during parsing
+     */
+    public static AggregationBuilder fromProto(Map.Entry<String, AggregationContainer> aggEntry) throws IOException {
+        String aggregationName = aggEntry.getKey();
+        AggregationContainer aggregationContainer = aggEntry.getValue();
+        switch (aggregationContainer.getAggregationCase()) {
+            case FILTER -> throw new IllegalArgumentException("Top level aggregation is a filter clause. Use the 'query' field to filter all results.");
+            case CARDINALITY -> {
+                return CardinalityAggregationBuilderProtoUtils.fromProto(aggregationContainer.getCardinality(), aggregationName);
+            }
+            case MISSING -> {
+                return MissingAggregationBuilderProtoUtils.fromProto(aggregationContainer.getMissing(), aggregationName);
+            }
+            case AGGREGATION_NOT_SET -> {
+                return null;
+            }
+        }
+        return null;
+    }
+}