Skip to content

Commit b6b9404

Browse files
committed
[GRPC] Add terms, Cardinality, Missing aggregations
Signed-off-by: Karen X <karenxyr@gmail.com>
1 parent 2f7e9c7 commit b6b9404

15 files changed

+676
-10
lines changed

gradle/libs.versions.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ kotlin = "1.7.10"
2222
antlr4 = "4.13.1"
2323
guava = "33.2.1-jre"
2424
gson = "2.13.2"
25-
opensearchprotobufs = "0.23.0"
25+
opensearchprotobufs = "0.24.0-SNAPSHOT"
2626
protobuf = "3.25.8"
2727
jakarta_annotation = "1.3.5"
2828
google_http_client = "1.44.1"

modules/transport-grpc/build.gradle

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,12 @@ dependencies {
3838
implementation "org.opensearch:protobufs:${versions.opensearchprotobufs}"
3939
testImplementation project(':test:framework')
4040
}
41+
repositories {
42+
maven {
43+
url = 'https://ci.opensearch.org/ci/dbc/snapshots/maven/'
44+
}
45+
// mavenLocal()
46+
}
4147

4248
tasks.named("dependencyLicenses").configure {
4349
mapping from: /grpc-.*/, to: 'grpc'

modules/transport-grpc/licenses/protobufs-0.23.0.jar.sha1

Lines changed: 0 additions & 1 deletion
This file was deleted.
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
20aae22bf8609cb6a963f6897c09384f8d5ecb78

modules/transport-grpc/spi/build.gradle

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,3 +41,9 @@ thirdPartyAudit {
4141
'com.google.common.util.concurrent.ListenableFuture'
4242
)
4343
}
44+
45+
repositories {
46+
maven {
47+
url = 'https://ci.opensearch.org/ci/dbc/snapshots/maven/'
48+
}
49+
}

modules/transport-grpc/spi/licenses/protobufs-0.23.0.jar.sha1

Lines changed: 0 additions & 1 deletion
This file was deleted.
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
20aae22bf8609cb6a963f6897c09384f8d5ecb78

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

Lines changed: 26 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -9,19 +9,23 @@
99

1010
import org.opensearch.common.unit.TimeValue;
1111
import org.opensearch.core.xcontent.XContentParser;
12+
import org.opensearch.protobufs.AggregationContainer;
1213
import org.opensearch.protobufs.DerivedField;
1314
import org.opensearch.protobufs.FieldAndFormat;
1415
import org.opensearch.protobufs.Rescore;
1516
import org.opensearch.protobufs.ScriptField;
1617
import org.opensearch.protobufs.SearchRequestBody;
1718
import org.opensearch.protobufs.TrackHits;
19+
import org.opensearch.search.aggregations.AggregationBuilder;
1820
import org.opensearch.search.builder.SearchSourceBuilder;
1921
import org.opensearch.search.sort.SortBuilder;
2022
import org.opensearch.transport.grpc.proto.request.common.FetchSourceContextProtoUtils;
2123
import org.opensearch.transport.grpc.proto.request.common.ScriptProtoUtils;
24+
import org.opensearch.transport.grpc.proto.request.search.aggregation.AggregationContainerProtoUtils;
2225
import org.opensearch.transport.grpc.proto.request.search.query.AbstractQueryBuilderProtoUtils;
2326
import org.opensearch.transport.grpc.proto.request.search.sort.SortBuilderProtoUtils;
2427
import org.opensearch.transport.grpc.proto.request.search.suggest.SuggestBuilderProtoUtils;
28+
import org.opensearch.transport.grpc.spi.QueryBuilderProtoConverterRegistry;
2529

2630
import java.io.IOException;
2731
import java.util.Map;
@@ -55,7 +59,7 @@ public static void parseProto(
5559
AbstractQueryBuilderProtoUtils queryUtils
5660
) throws IOException {
5761
// Parse all non-query fields
58-
parseNonQueryFields(searchSourceBuilder, protoRequest);
62+
parseNonQueryFields(searchSourceBuilder, protoRequest, queryUtils);
5963

6064
// Handle queries using the instance-based approach
6165
if (protoRequest.hasQuery()) {
@@ -67,9 +71,20 @@ public static void parseProto(
6771
}
6872

6973
/**
70-
* Parses all fields except queries from the protobuf SearchRequestBody.
74+
* Parses all fields from the protobuf SearchRequestBody in the same order as REST API.
75+
* This matches the order in {@link SearchSourceBuilder#parseXContent(XContentParser, boolean)}.
76+
*
77+
* @param searchSourceBuilder The SearchSourceBuilder to populate
78+
* @param protoRequest The Protocol Buffer SearchRequest to parse
79+
* @param queryUtils The query utils instance for parsing queries in aggregations/filters
80+
* @throws IOException if there's an error during parsing
7181
*/
72-
private static void parseNonQueryFields(SearchSourceBuilder searchSourceBuilder, SearchRequestBody protoRequest) throws IOException {
82+
private static void parseNonQueryFields(
83+
SearchSourceBuilder searchSourceBuilder,
84+
SearchRequestBody protoRequest,
85+
AbstractQueryBuilderProtoUtils queryUtils
86+
) throws IOException {
87+
QueryBuilderProtoConverterRegistry registry = queryUtils.getRegistry();
7388
// TODO what to do about parser.getDeprecationHandler() for protos?
7489

7590
if (protoRequest.hasFrom()) {
@@ -148,10 +163,14 @@ private static void parseNonQueryFields(SearchSourceBuilder searchSourceBuilder,
148163
}
149164
}
150165

151-
// TODO support aggregations
152-
/*
153-
if(protoRequest.hasAggs()){}
154-
*/
166+
if (protoRequest.getAggregationsCount() > 0) {
167+
for (Map.Entry<String, AggregationContainer> entry : protoRequest.getAggregationsMap().entrySet()) {
168+
String aggName = entry.getKey();
169+
AggregationContainer aggContainer = entry.getValue();
170+
AggregationBuilder aggBuilder = AggregationContainerProtoUtils.fromProto(aggName, aggContainer, queryUtils);
171+
searchSourceBuilder.aggregation(aggBuilder);
172+
}
173+
}
155174

156175
if (protoRequest.hasHighlight()) {
157176
searchSourceBuilder.highlighter(HighlightBuilderProtoUtils.fromProto(protoRequest.getHighlight()));
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,119 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
package org.opensearch.transport.grpc.proto.request.search.aggregation;
9+
10+
import org.opensearch.core.xcontent.XContentParser;
11+
import org.opensearch.protobufs.AggregationContainer;
12+
import org.opensearch.protobufs.ObjectMap;
13+
import org.opensearch.search.aggregations.AggregationBuilder;
14+
import org.opensearch.transport.grpc.proto.request.common.ObjectMapProtoUtils;
15+
import org.opensearch.transport.grpc.proto.request.search.query.AbstractQueryBuilderProtoUtils;
16+
17+
import java.util.Map;
18+
19+
/**
20+
* Utility class for converting AggregationContainer Protocol Buffers to OpenSearch AggregationBuilder objects.
21+
* This class handles the transformation of Protocol Buffer aggregation containers into their corresponding
22+
* OpenSearch AggregationBuilder implementations. It serves as the main entry point for converting all
23+
* aggregation types from protobuf format to OpenSearch format.
24+
*
25+
* <p>This class is analogous to the REST-side aggregation parsing framework where {@link XContentParser}
26+
* is used to parse aggregations from JSON/REST requests. Each aggregation type's {@code PARSER} field
27+
* (which uses {@code fromXContent}) has a corresponding protobuf converter method that performs the same
28+
* logical transformation but from Protocol Buffer representation.
29+
*
30+
* <p>The dispatch logic mirrors how REST aggregations are parsed: by examining the aggregation type field
31+
* and delegating to the appropriate type-specific parser.
32+
*/
33+
public class AggregationContainerProtoUtils {
34+
35+
private AggregationContainerProtoUtils() {
36+
// Utility class, no instances
37+
}
38+
39+
/**
40+
* Converts a Protocol Buffer AggregationContainer to an OpenSearch AggregationBuilder.
41+
* This method dispatches to the appropriate aggregation-specific converter based on the
42+
* aggregation type specified in the container.
43+
*
44+
* <p>This is the gRPC equivalent of the REST-side aggregation parsing in
45+
* {@code AggregationBuilder.parseAggregators}, where each aggregation type has a registered
46+
* parser that reads from {@link XContentParser} via {@code fromXContent}. This method performs
47+
* the same role but reads from Protocol Buffer structures instead of JSON.
48+
*
49+
* @param name The name of the aggregation
50+
* @param aggContainer The Protocol Buffer AggregationContainer object
51+
* @return A configured AggregationBuilder instance
52+
* @throws IllegalArgumentException if the aggregation type is not supported or unrecognized
53+
* @see org.opensearch.search.aggregations.AggregationBuilder
54+
*/
55+
public static AggregationBuilder fromProto(String name, AggregationContainer aggContainer) {
56+
return fromProto(name, aggContainer, null);
57+
}
58+
59+
/**
60+
* Converts a Protocol Buffer AggregationContainer to an OpenSearch AggregationBuilder.
61+
* This method dispatches to the appropriate aggregation-specific converter based on the
62+
* aggregation type specified in the container. It also supports nested queries in aggregations
63+
* (like filter aggregations) by accepting a queryUtils parameter.
64+
*
65+
* <p>This is the gRPC equivalent of the REST-side aggregation parsing in
66+
* {@code AggregationBuilder.parseAggregators}, where each aggregation type has a registered
67+
* parser that reads from {@link XContentParser} via {@code fromXContent}. This method performs
68+
* the same role but reads from Protocol Buffer structures instead of JSON.
69+
*
70+
* @param name The name of the aggregation
71+
* @param aggContainer The Protocol Buffer AggregationContainer object
72+
* @param queryUtils The query utils instance for parsing nested queries (can be null if not needed)
73+
* @return A configured AggregationBuilder instance
74+
* @throws IllegalArgumentException if the aggregation type is not supported or unrecognized
75+
* @see org.opensearch.search.aggregations.AggregationBuilder
76+
*/
77+
public static AggregationBuilder fromProto(String name, AggregationContainer aggContainer, AbstractQueryBuilderProtoUtils queryUtils) {
78+
AggregationBuilder builder;
79+
80+
// Dispatch based on the aggregation type
81+
switch (aggContainer.getAggregationCase()) {
82+
case CARDINALITY:
83+
builder = CardinalityAggregationProtoUtils.fromProto(name, aggContainer.getCardinality());
84+
break;
85+
86+
case MISSING:
87+
builder = MissingAggregationProtoUtils.fromProto(name, aggContainer.getMissing());
88+
break;
89+
90+
case TERMS:
91+
builder = TermsAggregationProtoUtils.fromProto(name, aggContainer.getTerms());
92+
break;
93+
94+
case FILTER:
95+
if (queryUtils == null) {
96+
throw new IllegalArgumentException("Filter aggregation requires queryUtils to be provided for parsing nested queries");
97+
}
98+
builder = FilterAggregationProtoUtils.fromProto(name, aggContainer.getFilter(), queryUtils);
99+
break;
100+
101+
case AGGREGATION_NOT_SET:
102+
throw new IllegalArgumentException("Aggregation type not set for aggregation: " + name);
103+
104+
default:
105+
throw new IllegalArgumentException(
106+
"Unsupported aggregation type: " + aggContainer.getAggregationCase() + " for aggregation: " + name
107+
);
108+
}
109+
110+
// Handle metadata if present
111+
if (aggContainer.hasMeta()) {
112+
ObjectMap metaProto = aggContainer.getMeta();
113+
Map<String, Object> metadata = ObjectMapProtoUtils.fromProto(metaProto);
114+
builder.setMetadata(metadata);
115+
}
116+
117+
return builder;
118+
}
119+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,111 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
package org.opensearch.transport.grpc.proto.request.search.aggregation;
9+
10+
import org.opensearch.protobufs.CardinalityAggregation;
11+
import org.opensearch.protobufs.CardinalityExecutionMode;
12+
import org.opensearch.protobufs.FieldValue;
13+
import org.opensearch.search.aggregations.metrics.CardinalityAggregationBuilder;
14+
import org.opensearch.transport.grpc.proto.request.common.ScriptProtoUtils;
15+
import org.opensearch.transport.grpc.proto.response.common.FieldValueProtoUtils;
16+
import org.opensearch.transport.grpc.util.ProtobufEnumUtils;
17+
18+
import java.lang.reflect.Method;
19+
20+
/**
21+
* Utility class for converting CardinalityAggregation Protocol Buffers to OpenSearch objects.
22+
* This class provides methods to transform Protocol Buffer representations of cardinality aggregations
23+
* into their corresponding OpenSearch CardinalityAggregationBuilder implementations.
24+
*/
25+
public class CardinalityAggregationProtoUtils {
26+
27+
private CardinalityAggregationProtoUtils() {
28+
// Utility class, no instances
29+
}
30+
31+
/**
32+
* Converts a Protocol Buffer CardinalityAggregation to an OpenSearch CardinalityAggregationBuilder.
33+
*
34+
* <p>This method is the gRPC equivalent of {@link CardinalityAggregationBuilder#PARSER}, which parses
35+
* cardinality aggregations from REST/JSON requests via {@code fromXContent}. Similar to how the parser
36+
* reads JSON fields, this method extracts values from the Protocol Buffer representation and creates
37+
* a properly configured CardinalityAggregationBuilder.
38+
*
39+
* <p>The REST-side serialization via {@link CardinalityAggregationBuilder#doXContentBody} produces
40+
* JSON that conceptually mirrors the protobuf structure used here.
41+
*
42+
* @param name The name of the aggregation
43+
* @param cardinalityAggProto The Protocol Buffer CardinalityAggregation object
44+
* @return A configured CardinalityAggregationBuilder instance
45+
* @throws IllegalArgumentException if the field value type is not supported
46+
* @see CardinalityAggregationBuilder#PARSER
47+
* @see CardinalityAggregationBuilder#doXContentBody
48+
*/
49+
public static CardinalityAggregationBuilder fromProto(String name, CardinalityAggregation cardinalityAggProto) {
50+
CardinalityAggregationBuilder builder = new CardinalityAggregationBuilder(name);
51+
52+
// Set field if present
53+
if (cardinalityAggProto.hasField()) {
54+
builder.field(cardinalityAggProto.getField());
55+
}
56+
57+
// Set missing value if present
58+
if (cardinalityAggProto.hasMissing()) {
59+
FieldValue missingValue = cardinalityAggProto.getMissing();
60+
Object missing = FieldValueProtoUtils.fromProto(missingValue, false);
61+
builder.missing(missing);
62+
}
63+
64+
// Set script if present
65+
if (cardinalityAggProto.hasScript()) {
66+
try {
67+
builder.script(ScriptProtoUtils.parseFromProtoRequest(cardinalityAggProto.getScript()));
68+
} catch (Exception e) {
69+
throw new IllegalArgumentException("Failed to parse script for cardinality aggregation", e);
70+
}
71+
}
72+
73+
// Set precision threshold if present
74+
if (cardinalityAggProto.hasPrecisionThreshold()) {
75+
builder.precisionThreshold(cardinalityAggProto.getPrecisionThreshold());
76+
}
77+
78+
// Set execution hint if present (added in OpenSearch 2.19.1)
79+
// Use reflection for forward-compatibility: if the OpenSearch version doesn't support
80+
// this field yet, it will be silently ignored rather than causing a compilation error.
81+
if (cardinalityAggProto.hasExecutionHint()) {
82+
String executionHint = convertExecutionMode(cardinalityAggProto.getExecutionHint());
83+
if (executionHint != null) {
84+
try {
85+
Method method = CardinalityAggregationBuilder.class.getMethod("executionHint", String.class);
86+
method.invoke(builder, executionHint);
87+
} catch (NoSuchMethodException e) {
88+
// Method doesn't exist in OpenSearch < 2.19.1 - silently ignore
89+
} catch (Exception e) {
90+
throw new IllegalArgumentException("Failed to set execution hint for cardinality aggregation", e);
91+
}
92+
}
93+
}
94+
95+
return builder;
96+
}
97+
98+
/**
99+
* Converts a Protocol Buffer CardinalityExecutionMode to its string representation.
100+
* Supports execution hints added in OpenSearch 2.19.1.
101+
*
102+
* @param mode The Protocol Buffer CardinalityExecutionMode
103+
* @return The string representation of the execution mode, or null if unspecified
104+
*/
105+
private static String convertExecutionMode(CardinalityExecutionMode mode) {
106+
if (mode == CardinalityExecutionMode.CARDINALITY_EXECUTION_MODE_UNSPECIFIED) {
107+
return null;
108+
}
109+
return ProtobufEnumUtils.convertToString(mode);
110+
}
111+
}

0 commit comments

Comments
 (0)