Skip to content

Commit 4e0aa21

Browse files
VachaShahVachaShah
committed
Adding protos for search classes
Signed-off-by: Vacha Shah <vachshah@amazon.com>
1 parent 14f1c43 commit 4e0aa21

27 files changed

+1349
-3
lines changed

‎libs/common/src/main/java/org/opensearch/common/annotation/processor/ApiAnnotationProcessor.java

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -238,7 +238,15 @@ private boolean inspectable(ExecutableElement executable) {
238238
*/
239239
private boolean inspectable(Element element) {
240240
final PackageElement pckg = processingEnv.getElementUtils().getPackageOf(element);
241-
return pckg.getQualifiedName().toString().startsWith(OPENSEARCH_PACKAGE);
241+
return pckg.getQualifiedName().toString().startsWith(OPENSEARCH_PACKAGE)
242+
&& !element.getEnclosingElement()
243+
.getAnnotationMirrors()
244+
.stream()
245+
.anyMatch(
246+
m -> m.getAnnotationType()
247+
.toString() /* ClassSymbol.toString() returns class name */
248+
.equalsIgnoreCase("javax.annotation.Generated")
249+
);
242250
}
243251

244252
/**

‎server/build.gradle

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -366,6 +366,32 @@ tasks.named("dependencyLicenses").configure {
366366
}
367367
}
368368

369+
tasks.named("missingJavadoc").configure {
370+
/*
371+
* annotate_code in L210 does not add the Generated annotation to nested code generated using protobuf.
372+
* TODO: Add support to missingJavadoc task to ignore all such nested classes.
373+
* https://github.com/opensearch-project/OpenSearch/issues/11913
374+
*/
375+
dependsOn("generateProto")
376+
javadocMissingIgnore = [
377+
"org.opensearch.server.proto.QuerySearchResultProto.QuerySearchResult.RescoreDocIds.setIntegerOrBuilder",
378+
"org.opensearch.server.proto.QuerySearchResultProto.QuerySearchResult.RescoreDocIdsOrBuilder",
379+
"org.opensearch.server.proto.QuerySearchResultProto.QuerySearchResult.TopDocs.ScoreDocOrBuilder",
380+
"org.opensearch.server.proto.QuerySearchResultProto.QuerySearchResult.TopDocsOrBuilder",
381+
"org.opensearch.server.proto.QuerySearchResultProto.QuerySearchResult.TopDocsAndMaxScoreOrBuilder",
382+
"org.opensearch.server.proto.FetchSearchResultProto.SearchHit.SearchSortValuesOrBuilder",
383+
"org.opensearch.server.proto.FetchSearchResultProto.SearchHit.HighlightFieldOrBuilder",
384+
"org.opensearch.server.proto.FetchSearchResultProto.SearchHit.DocumentFieldOrBuilder",
385+
"org.opensearch.server.proto.FetchSearchResultProto.SearchHit.NestedIdentityOrBuilder",
386+
"org.opensearch.server.proto.NodeToNodeMessageProto.NodeToNodeMessage.MessageCase",
387+
"org.opensearch.server.proto.NodeToNodeMessageProto.NodeToNodeMessage.ResponseHandlersListOrBuilder",
388+
"org.opensearch.server.proto.NodeToNodeMessageProto.NodeToNodeMessage.HeaderOrBuilder",
389+
"org.opensearch.server.proto.FetchSearchResultProto.SearchHit.Explanation.ExplanationValueCase",
390+
"org.opensearch.server.proto.FetchSearchResultProto.SearchHit.ExplanationOrBuilder",
391+
"org.opensearch.server.proto.ShardSearchRequestProto.OriginalIndices.IndicesOptionsOrBuilder",
392+
]
393+
}
394+
369395
tasks.named("filepermissions").configure {
370396
mustRunAfter("generateProto")
371397
}
@@ -380,6 +406,7 @@ tasks.named("licenseHeaders").configure {
380406
excludes << 'org/opensearch/client/documentation/placeholder.txt'
381407
// Ignore for protobuf generated code
382408
excludes << 'org/opensearch/extensions/proto/*'
409+
excludes << 'org/opensearch/server/proto/*'
383410
}
384411

385412
tasks.test {

‎server/src/main/java/org/opensearch/common/document/serializer/DocumentFieldProtobufSerializer.java

Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,143 @@
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+
9+
package org.opensearch.common.document.serializer;
10+
11+
import com.google.protobuf.ByteString;
12+
import org.opensearch.OpenSearchException;
13+
import org.opensearch.common.document.DocumentField;
14+
import org.opensearch.core.common.text.Text;
15+
import org.opensearch.server.proto.FetchSearchResultProto;
16+
import org.opensearch.server.proto.FetchSearchResultProto.DocumentFieldValue;
17+
import org.opensearch.server.proto.FetchSearchResultProto.DocumentFieldValue.Builder;
18+
19+
import java.io.IOException;
20+
import java.io.InputStream;
21+
import java.time.Instant;
22+
import java.time.ZoneId;
23+
import java.time.ZonedDateTime;
24+
import java.util.ArrayList;
25+
import java.util.Date;
26+
import java.util.HashMap;
27+
import java.util.LinkedHashMap;
28+
import java.util.List;
29+
import java.util.Map;
30+
31+
/**
32+
* Serializer for {@link DocumentField} to/from protobuf.
33+
*/
34+
public class DocumentFieldProtobufSerializer implements DocumentFieldSerializer<InputStream> {
35+
36+
private FetchSearchResultProto.SearchHit.DocumentField documentField;
37+
38+
@Override
39+
public DocumentField createDocumentField(InputStream inputStream) throws IOException {
40+
documentField = FetchSearchResultProto.SearchHit.DocumentField.parseFrom(inputStream);
41+
String name = documentField.getName();
42+
List<Object> values = new ArrayList<>();
43+
for (FetchSearchResultProto.DocumentFieldValue value : documentField.getValuesList()) {
44+
values.add(readDocumentFieldValueFromProtobuf(value));
45+
}
46+
return new DocumentField(name, values);
47+
}
48+
49+
private Object readDocumentFieldValueFromProtobuf(FetchSearchResultProto.DocumentFieldValue documentFieldValue) throws IOException {
50+
if (documentFieldValue.hasValueString()) {
51+
return documentFieldValue.getValueString();
52+
} else if (documentFieldValue.hasValueInt()) {
53+
return documentFieldValue.getValueInt();
54+
} else if (documentFieldValue.hasValueLong()) {
55+
return documentFieldValue.getValueLong();
56+
} else if (documentFieldValue.hasValueFloat()) {
57+
return documentFieldValue.getValueFloat();
58+
} else if (documentFieldValue.hasValueDouble()) {
59+
return documentFieldValue.getValueDouble();
60+
} else if (documentFieldValue.hasValueBool()) {
61+
return documentFieldValue.getValueBool();
62+
} else if (documentFieldValue.getValueByteArrayList().size() > 0) {
63+
return documentFieldValue.getValueByteArrayList().toArray();
64+
} else if (documentFieldValue.getValueArrayListList().size() > 0) {
65+
List<Object> list = new ArrayList<>();
66+
for (FetchSearchResultProto.DocumentFieldValue value : documentFieldValue.getValueArrayListList()) {
67+
list.add(readDocumentFieldValueFromProtobuf(value));
68+
}
69+
return list;
70+
} else if (documentFieldValue.getValueMapMap().size() > 0) {
71+
Map<String, Object> map = Map.of();
72+
for (Map.Entry<String, FetchSearchResultProto.DocumentFieldValue> entrySet : documentFieldValue.getValueMapMap().entrySet()) {
73+
map.put(entrySet.getKey(), readDocumentFieldValueFromProtobuf(entrySet.getValue()));
74+
}
75+
return map;
76+
} else if (documentFieldValue.hasValueDate()) {
77+
return new Date(documentFieldValue.getValueDate());
78+
} else if (documentFieldValue.hasValueZonedDate() && documentFieldValue.hasValueZonedTime()) {
79+
return ZonedDateTime.ofInstant(
80+
Instant.ofEpochMilli(documentFieldValue.getValueZonedTime()),
81+
ZoneId.of(documentFieldValue.getValueZonedDate())
82+
);
83+
} else if (documentFieldValue.hasValueText()) {
84+
return new Text(documentFieldValue.getValueText());
85+
} else {
86+
throw new IOException("Can't read generic value of type [" + documentFieldValue + "]");
87+
}
88+
}
89+
90+
public static DocumentFieldValue.Builder convertDocumentFieldValueToProto(Object value, Builder valueBuilder) {
91+
if (value == null) {
92+
// null is not allowed in protobuf, so we use a special string to represent null
93+
return valueBuilder.setValueString("null");
94+
}
95+
Class type = value.getClass();
96+
if (type == String.class) {
97+
valueBuilder.setValueString((String) value);
98+
} else if (type == Integer.class) {
99+
valueBuilder.setValueInt((Integer) value);
100+
} else if (type == Long.class) {
101+
valueBuilder.setValueLong((Long) value);
102+
} else if (type == Float.class) {
103+
valueBuilder.setValueFloat((Float) value);
104+
} else if (type == Double.class) {
105+
valueBuilder.setValueDouble((Double) value);
106+
} else if (type == Boolean.class) {
107+
valueBuilder.setValueBool((Boolean) value);
108+
} else if (type == byte[].class) {
109+
valueBuilder.addValueByteArray(ByteString.copyFrom((byte[]) value));
110+
} else if (type == List.class) {
111+
List<Object> list = (List<Object>) value;
112+
for (Object listValue : list) {
113+
valueBuilder.addValueArrayList(convertDocumentFieldValueToProto(listValue, valueBuilder));
114+
}
115+
} else if (type == Map.class || type == HashMap.class || type == LinkedHashMap.class) {
116+
Map<String, Object> map = (Map<String, Object>) value;
117+
for (Map.Entry<String, Object> entry : map.entrySet()) {
118+
valueBuilder.putValueMap(entry.getKey(), convertDocumentFieldValueToProto(entry.getValue(), valueBuilder).build());
119+
}
120+
} else if (type == Date.class) {
121+
valueBuilder.setValueDate(((Date) value).getTime());
122+
} else if (type == ZonedDateTime.class) {
123+
valueBuilder.setValueZonedDate(((ZonedDateTime) value).getZone().getId());
124+
valueBuilder.setValueZonedTime(((ZonedDateTime) value).toInstant().toEpochMilli());
125+
} else if (type == Text.class) {
126+
valueBuilder.setValueText(((Text) value).string());
127+
} else {
128+
throw new OpenSearchException("Can't convert generic value of type [" + type + "] to protobuf");
129+
}
130+
return valueBuilder;
131+
}
132+
133+
public static FetchSearchResultProto.SearchHit.DocumentField convertDocumentFieldToProto(DocumentField documentField) {
134+
FetchSearchResultProto.SearchHit.DocumentField.Builder builder = FetchSearchResultProto.SearchHit.DocumentField.newBuilder();
135+
builder.setName(documentField.getName());
136+
for (Object value : documentField.getValues()) {
137+
FetchSearchResultProto.DocumentFieldValue.Builder valueBuilder = FetchSearchResultProto.DocumentFieldValue.newBuilder();
138+
builder.addValues(DocumentFieldProtobufSerializer.convertDocumentFieldValueToProto(value, valueBuilder));
139+
}
140+
return builder.build();
141+
}
142+
143+
}

‎server/src/main/java/org/opensearch/common/document/serializer/DocumentFieldSerializer.java

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
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+
9+
package org.opensearch.common.document.serializer;
10+
11+
import org.opensearch.common.document.DocumentField;
12+
13+
import java.io.IOException;
14+
15+
/**
16+
* Serializer for {@link DocumentField} which can be implemented for different types of serialization.
17+
*/
18+
public interface DocumentFieldSerializer<T> {
19+
20+
DocumentField createDocumentField(T inputStream) throws IOException;
21+
22+
}

‎server/src/main/java/org/opensearch/common/document/serializer/package-info.java

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
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+
9+
/** Serializer package for documents. */
10+
package org.opensearch.common.document.serializer;

‎server/src/main/java/org/opensearch/common/util/FeatureFlags.java

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,11 @@ public class FeatureFlags {
6767
*/
6868
public static final String PLUGGABLE_CACHE = "opensearch.experimental.feature.pluggable.caching.enabled";
6969

70+
/**
71+
* Gates the functionality of integrating protobuf within search API and node-to-node communication.
72+
*/
73+
public static final String PROTOBUF = "opensearch.experimental.feature.search_with_protobuf.enabled";
74+
7075
public static final Setting<Boolean> REMOTE_STORE_MIGRATION_EXPERIMENTAL_SETTING = Setting.boolSetting(
7176
REMOTE_STORE_MIGRATION_EXPERIMENTAL,
7277
false,
@@ -93,14 +98,17 @@ public class FeatureFlags {
9398

9499
public static final Setting<Boolean> PLUGGABLE_CACHE_SETTING = Setting.boolSetting(PLUGGABLE_CACHE, false, Property.NodeScope);
95100

101+
public static final Setting<Boolean> PROTOBUF_SETTING = Setting.boolSetting(PROTOBUF, false, Property.NodeScope, Property.Dynamic);
102+
96103
private static final List<Setting<Boolean>> ALL_FEATURE_FLAG_SETTINGS = List.of(
97104
REMOTE_STORE_MIGRATION_EXPERIMENTAL_SETTING,
98105
EXTENSIONS_SETTING,
99106
IDENTITY_SETTING,
100107
TELEMETRY_SETTING,
101108
DATETIME_FORMATTER_CACHING_SETTING,
102109
WRITEABLE_REMOTE_INDEX_SETTING,
103-
PLUGGABLE_CACHE_SETTING
110+
PLUGGABLE_CACHE_SETTING,
111+
PROTOBUF_SETTING
104112
);
105113
/**
106114
* Should store the settings from opensearch.yml.

‎server/src/main/java/org/opensearch/search/SearchSortValues.java

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,11 @@ public class SearchSortValues implements ToXContentFragment, Writeable {
6767
this.rawSortValues = EMPTY_ARRAY;
6868
}
6969

70+
public SearchSortValues(Object[] sortValues, Object[] rawSortValues) {
71+
this.formattedSortValues = Objects.requireNonNull(sortValues, "sort values must not be empty");
72+
this.rawSortValues = rawSortValues;
73+
}
74+
7075
public SearchSortValues(Object[] rawSortValues, DocValueFormat[] sortValueFormats) {
7176
Objects.requireNonNull(rawSortValues);
7277
Objects.requireNonNull(sortValueFormats);

‎server/src/main/java/org/opensearch/search/fetch/FetchSearchResult.java

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
package org.opensearch.search.fetch;
3434

3535
import org.opensearch.common.annotation.PublicApi;
36+
import org.opensearch.common.util.FeatureFlags;
3637
import org.opensearch.core.common.io.stream.StreamInput;
3738
import org.opensearch.core.common.io.stream.StreamOutput;
3839
import org.opensearch.search.SearchHit;
@@ -41,8 +42,13 @@
4142
import org.opensearch.search.SearchShardTarget;
4243
import org.opensearch.search.internal.ShardSearchContextId;
4344
import org.opensearch.search.query.QuerySearchResult;
45+
import org.opensearch.search.serializer.SearchHitsProtobufSerializer;
46+
import org.opensearch.server.proto.FetchSearchResultProto;
47+
import org.opensearch.server.proto.ShardSearchRequestProto;
4448

49+
import java.io.ByteArrayInputStream;
4550
import java.io.IOException;
51+
import java.io.InputStream;
4652

4753
/**
4854
* Result from a fetch
@@ -56,6 +62,8 @@ public final class FetchSearchResult extends SearchPhaseResult {
5662
// client side counter
5763
private transient int counter;
5864

65+
private FetchSearchResultProto.FetchSearchResult fetchSearchResultProto;
66+
5967
public FetchSearchResult() {}
6068

6169
public FetchSearchResult(StreamInput in) throws IOException {
@@ -64,9 +72,24 @@ public FetchSearchResult(StreamInput in) throws IOException {
6472
hits = new SearchHits(in);
6573
}
6674

75+
public FetchSearchResult(InputStream in) throws IOException {
76+
this.fetchSearchResultProto = FetchSearchResultProto.FetchSearchResult.parseFrom(in);
77+
contextId = new ShardSearchContextId(
78+
this.fetchSearchResultProto.getContextId().getSessionId(),
79+
this.fetchSearchResultProto.getContextId().getId()
80+
);
81+
SearchHitsProtobufSerializer protobufSerializer = new SearchHitsProtobufSerializer();
82+
hits = protobufSerializer.createSearchHits(new ByteArrayInputStream(this.fetchSearchResultProto.getHits().toByteArray()));
83+
}
84+
6785
public FetchSearchResult(ShardSearchContextId id, SearchShardTarget shardTarget) {
6886
this.contextId = id;
6987
setSearchShardTarget(shardTarget);
88+
this.fetchSearchResultProto = FetchSearchResultProto.FetchSearchResult.newBuilder()
89+
.setContextId(
90+
ShardSearchRequestProto.ShardSearchContextId.newBuilder().setSessionId(id.getSessionId()).setId(id.getId()).build()
91+
)
92+
.build();
7093
}
7194

7295
@Override
@@ -82,6 +105,11 @@ public FetchSearchResult fetchResult() {
82105
public void hits(SearchHits hits) {
83106
assert assertNoSearchTarget(hits);
84107
this.hits = hits;
108+
if (FeatureFlags.isEnabled(FeatureFlags.PROTOBUF_SETTING) && this.fetchSearchResultProto != null) {
109+
this.fetchSearchResultProto = this.fetchSearchResultProto.toBuilder()
110+
.setHits(SearchHitsProtobufSerializer.convertHitsToProto(hits))
111+
.build();
112+
}
85113
}
86114

87115
private boolean assertNoSearchTarget(SearchHits hits) {
@@ -92,6 +120,16 @@ private boolean assertNoSearchTarget(SearchHits hits) {
92120
}
93121

94122
public SearchHits hits() {
123+
if (FeatureFlags.isEnabled(FeatureFlags.PROTOBUF_SETTING) && this.fetchSearchResultProto != null) {
124+
SearchHits hits;
125+
try {
126+
SearchHitsProtobufSerializer protobufSerializer = new SearchHitsProtobufSerializer();
127+
hits = protobufSerializer.createSearchHits(new ByteArrayInputStream(this.fetchSearchResultProto.getHits().toByteArray()));
128+
return hits;
129+
} catch (IOException e) {
130+
throw new RuntimeException(e);
131+
}
132+
}
95133
return hits;
96134
}
97135

‎server/src/main/java/org/opensearch/search/fetch/QueryFetchSearchResult.java

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,8 +38,10 @@
3838
import org.opensearch.search.SearchShardTarget;
3939
import org.opensearch.search.internal.ShardSearchContextId;
4040
import org.opensearch.search.query.QuerySearchResult;
41+
import org.opensearch.server.proto.QueryFetchSearchResultProto;
4142

4243
import java.io.IOException;
44+
import java.io.InputStream;
4345

4446
/**
4547
* Query fetch result
@@ -51,12 +53,20 @@ public final class QueryFetchSearchResult extends SearchPhaseResult {
5153
private final QuerySearchResult queryResult;
5254
private final FetchSearchResult fetchResult;
5355

56+
private QueryFetchSearchResultProto.QueryFetchSearchResult queryFetchSearchResultProto;
57+
5458
public QueryFetchSearchResult(StreamInput in) throws IOException {
5559
super(in);
5660
queryResult = new QuerySearchResult(in);
5761
fetchResult = new FetchSearchResult(in);
5862
}
5963

64+
public QueryFetchSearchResult(InputStream in) throws IOException {
65+
this.queryFetchSearchResultProto = QueryFetchSearchResultProto.QueryFetchSearchResult.parseFrom(in);
66+
queryResult = new QuerySearchResult(in);
67+
fetchResult = new FetchSearchResult(in);
68+
}
69+
6070
public QueryFetchSearchResult(QuerySearchResult queryResult, FetchSearchResult fetchResult) {
6171
this.queryResult = queryResult;
6272
this.fetchResult = fetchResult;

0 commit comments

Comments
 (0)