Skip to content

Commit 83be522

Browse files
VachaShahVachaShah
committed
Adding protos for the remaining structures
Signed-off-by: Vacha Shah <vachshah@amazon.com>
1 parent abc32a8 commit 83be522

File tree

8 files changed

+267
-57
lines changed

8 files changed

+267
-57
lines changed

‎server/build.gradle

+3-1
Original file line numberDiff line numberDiff line change
@@ -368,7 +368,9 @@ tasks.named("missingJavadoc").configure {
368368
"org.opensearch.server.proto.FetchSearchResultProto.SearchHit.NestedIdentityOrBuilder",
369369
"org.opensearch.server.proto.NodeToNodeMessageProto.NodeToNodeMessage.MessageCase",
370370
"org.opensearch.server.proto.NodeToNodeMessageProto.NodeToNodeMessage.ResponseHandlersListOrBuilder",
371-
"org.opensearch.server.proto.NodeToNodeMessageProto.NodeToNodeMessage.HeaderOrBuilder"
371+
"org.opensearch.server.proto.NodeToNodeMessageProto.NodeToNodeMessage.HeaderOrBuilder",
372+
"org.opensearch.server.proto.FetchSearchResultProto.SearchHit.Explanation.ExplanationValueCase",
373+
"org.opensearch.server.proto.FetchSearchResultProto.SearchHit.ExplanationOrBuilder"
372374
]
373375
}
374376

‎server/src/main/java/org/opensearch/common/document/DocumentField.java

+23
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,9 @@
3232

3333
package org.opensearch.common.document;
3434

35+
import com.google.protobuf.ByteString;
36+
import org.apache.lucene.util.SuppressForbidden;
37+
import org.opensearch.OpenSearchException;
3538
import org.opensearch.common.annotation.PublicApi;
3639
import org.opensearch.core.common.io.stream.StreamInput;
3740
import org.opensearch.core.common.io.stream.StreamOutput;
@@ -41,8 +44,12 @@
4144
import org.opensearch.core.xcontent.XContentParser;
4245
import org.opensearch.index.get.GetResult;
4346
import org.opensearch.search.SearchHit;
47+
import org.opensearch.server.proto.FetchSearchResultProto;
4448

49+
import java.io.ByteArrayInputStream;
4550
import java.io.IOException;
51+
import java.io.InputStream;
52+
import java.io.ObjectInputStream;
4653
import java.util.ArrayList;
4754
import java.util.Iterator;
4855
import java.util.List;
@@ -64,6 +71,7 @@ public class DocumentField implements Writeable, ToXContentFragment, Iterable<Ob
6471

6572
private final String name;
6673
private final List<Object> values;
74+
private FetchSearchResultProto.SearchHit.DocumentField documentField;
6775

6876
public DocumentField(StreamInput in) throws IOException {
6977
name = in.readString();
@@ -75,6 +83,21 @@ public DocumentField(String name, List<Object> values) {
7583
this.values = Objects.requireNonNull(values, "values must not be null");
7684
}
7785

86+
@SuppressForbidden(reason = "We need to read from a byte array")
87+
public DocumentField(byte[] in) throws IOException {
88+
documentField = FetchSearchResultProto.SearchHit.DocumentField.parseFrom(in);
89+
name = documentField.getName();
90+
values = new ArrayList<>();
91+
for (ByteString value : documentField.getValuesList()) {
92+
InputStream is = new ByteArrayInputStream(value.toByteArray());
93+
try (ObjectInputStream ois = new ObjectInputStream(is)) {
94+
values.add(ois.readObject());
95+
} catch (ClassNotFoundException e) {
96+
throw new OpenSearchException(e);
97+
}
98+
}
99+
}
100+
78101
/**
79102
* The name of the field.
80103
*/

‎server/src/main/java/org/opensearch/common/lucene/Lucene.java

+24
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,7 @@
9393
import org.opensearch.index.analysis.NamedAnalyzer;
9494
import org.opensearch.index.fielddata.IndexFieldData;
9595
import org.opensearch.search.sort.SortedWiderNumericSortField;
96+
import org.opensearch.server.proto.FetchSearchResultProto;
9697

9798
import java.io.IOException;
9899
import java.math.BigInteger;
@@ -651,6 +652,29 @@ public static Explanation readExplanation(StreamInput in) throws IOException {
651652
}
652653
}
653654

655+
public static Explanation readExplanation(byte[] in) throws IOException {
656+
FetchSearchResultProto.SearchHit.Explanation explanationProto = FetchSearchResultProto.SearchHit.Explanation.parseFrom(in);
657+
boolean match = explanationProto.getMatch();
658+
String description = explanationProto.getDescription();
659+
final Explanation[] subExplanations = new Explanation[explanationProto.getSubExplanationsCount()];
660+
for (int i = 0; i < subExplanations.length; ++i) {
661+
subExplanations[i] = readExplanation(explanationProto.getSubExplanations(i).toByteArray());
662+
}
663+
Number explanationValue = null;
664+
if (explanationProto.hasValue1()) {
665+
explanationValue = explanationProto.getValue1();
666+
} else if (explanationProto.hasValue2()) {
667+
explanationValue = explanationProto.getValue2();
668+
} else if (explanationProto.hasValue3()) {
669+
explanationValue = explanationProto.getValue3();
670+
}
671+
if (match) {
672+
return Explanation.match(explanationValue, description, subExplanations);
673+
} else {
674+
return Explanation.noMatch(description, subExplanations);
675+
}
676+
}
677+
654678
private static void writeExplanationValue(StreamOutput out, Number value) throws IOException {
655679
if (value instanceof Float) {
656680
out.writeByte((byte) 0);

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

+75-31
Original file line numberDiff line numberDiff line change
@@ -234,48 +234,73 @@ public SearchHit(byte[] in) throws IOException {
234234
this.docId = -1;
235235
this.score = this.searchHitProto.getScore();
236236
this.id = new Text(this.searchHitProto.getId());
237-
// Support for nestedIdentity to be added in the future
238-
this.nestedIdentity = null;
237+
if (!this.searchHitProto.hasNestedIdentity() && this.searchHitProto.getNestedIdentity().toByteArray().length > 0) {
238+
this.nestedIdentity = new NestedIdentity(this.searchHitProto.getNestedIdentity().toByteArray());
239+
} else {
240+
this.nestedIdentity = null;
241+
}
239242
this.version = this.searchHitProto.getVersion();
240243
this.seqNo = this.searchHitProto.getSeqNo();
241244
this.primaryTerm = this.searchHitProto.getPrimaryTerm();
242245
this.source = BytesReference.fromByteBuffer(ByteBuffer.wrap(this.searchHitProto.getSource().toByteArray()));
243246
if (source.length() == 0) {
244247
source = null;
245248
}
246-
// add support for metaFields and documentFields
247-
this.metaFields = new HashMap<>();
248249
this.documentFields = new HashMap<>();
249-
}
250-
251-
private Map<String, DocumentField> readFields(StreamInput in) throws IOException {
252-
Map<String, DocumentField> fields;
253-
int size = in.readVInt();
254-
if (size == 0) {
255-
fields = emptyMap();
256-
} else if (size == 1) {
257-
DocumentField hitField = new DocumentField(in);
258-
fields = singletonMap(hitField.getName(), hitField);
259-
} else {
260-
fields = new HashMap<>(size);
261-
for (int i = 0; i < size; i++) {
262-
DocumentField field = new DocumentField(in);
263-
fields.put(field.getName(), field);
250+
this.searchHitProto.getDocumentFieldsMap().forEach((k, v) -> {
251+
try {
252+
this.documentFields.put(k, new DocumentField(v.toByteArray()));
253+
} catch (IOException e) {
254+
throw new OpenSearchParseException("failed to parse document field", e);
264255
}
265-
fields = unmodifiableMap(fields);
266-
}
267-
return fields;
268-
}
269-
270-
private void writeFields(StreamOutput out, Map<String, DocumentField> fields) throws IOException {
271-
if (fields == null) {
272-
out.writeVInt(0);
273-
} else {
274-
out.writeVInt(fields.size());
275-
for (DocumentField field : fields.values()) {
276-
field.writeTo(out);
256+
});
257+
this.metaFields = new HashMap<>();
258+
this.searchHitProto.getMetaFieldsMap().forEach((k, v) -> {
259+
try {
260+
this.metaFields.put(k, new DocumentField(v.toByteArray()));
261+
} catch (IOException e) {
262+
throw new OpenSearchParseException("failed to parse document field", e);
277263
}
264+
});
265+
this.highlightFields = new HashMap<>();
266+
this.searchHitProto.getHighlightFieldsMap().forEach((k, v) -> {
267+
try {
268+
this.highlightFields.put(k, new HighlightField(v.toByteArray()));
269+
} catch (IOException e) {
270+
throw new OpenSearchParseException("failed to parse highlight field", e);
271+
}
272+
});
273+
this.sortValues = new SearchSortValues(this.searchHitProto.getSortValues().toByteArray());
274+
if (this.searchHitProto.getMatchedQueriesCount() > 0) {
275+
this.matchedQueries = this.searchHitProto.getMatchedQueriesList().toArray(new String[0]);
276+
}
277+
if (this.searchHitProto.hasExplanation()) {
278+
this.explanation = readExplanation(this.searchHitProto.getExplanation().toByteArray());
279+
}
280+
SearchShardTarget searchShardTarget = new SearchShardTarget(
281+
this.searchHitProto.getShard().getNodeId(),
282+
new ShardId(
283+
this.searchHitProto.getShard().getShardId().getIndexName(),
284+
this.searchHitProto.getShard().getShardId().getIndexUUID(),
285+
this.searchHitProto.getShard().getShardId().getShardId()
286+
),
287+
this.searchHitProto.getShard().getClusterAlias(),
288+
OriginalIndices.NONE
289+
);
290+
shard(searchShardTarget);
291+
if (this.searchHitProto.getInnerHitsCount() > 0) {
292+
this.innerHits = new HashMap<>();
293+
this.searchHitProto.getInnerHitsMap().forEach((k, v) -> {
294+
try {
295+
this.innerHits.put(k, new SearchHits(v.toByteArray()));
296+
} catch (IOException e) {
297+
throw new OpenSearchParseException("failed to parse inner hits", e);
298+
}
299+
});
300+
} else {
301+
this.innerHits = null;
278302
}
303+
279304
}
280305

281306
private static final Text SINGLE_MAPPING_TYPE = new Text(MapperService.SINGLE_MAPPING_NAME);
@@ -1050,6 +1075,25 @@ public NestedIdentity(String field, int offset, NestedIdentity child) {
10501075
child = in.readOptionalWriteable(NestedIdentity::new);
10511076
}
10521077

1078+
NestedIdentity(byte[] in) throws IOException {
1079+
FetchSearchResultProto.SearchHit.NestedIdentity proto = FetchSearchResultProto.SearchHit.NestedIdentity.parseFrom(in);
1080+
if (proto.hasField()) {
1081+
field = new Text(proto.getField());
1082+
} else {
1083+
field = null;
1084+
}
1085+
if (proto.hasOffset()) {
1086+
offset = proto.getOffset();
1087+
} else {
1088+
offset = -1;
1089+
}
1090+
if (proto.hasChild()) {
1091+
child = new NestedIdentity(proto.getChild().toByteArray());
1092+
} else {
1093+
child = null;
1094+
}
1095+
}
1096+
10531097
/**
10541098
* Returns the nested field in the source this hit originates from
10551099
*/

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

+50-13
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,8 @@
3636
import org.apache.lucene.search.SortField;
3737
import org.apache.lucene.search.TotalHits;
3838
import org.apache.lucene.search.TotalHits.Relation;
39+
import org.apache.lucene.util.SuppressForbidden;
40+
import org.opensearch.OpenSearchException;
3941
import org.opensearch.common.Nullable;
4042
import org.opensearch.common.annotation.PublicApi;
4143
import org.opensearch.common.lucene.Lucene;
@@ -50,7 +52,12 @@
5052
import org.opensearch.server.proto.FetchSearchResultProto;
5153
import org.opensearch.server.proto.QuerySearchResultProto;
5254

55+
import java.io.ByteArrayInputStream;
56+
import java.io.ByteArrayOutputStream;
5357
import java.io.IOException;
58+
import java.io.InputStream;
59+
import java.io.ObjectInputStream;
60+
import java.io.ObjectOutputStream;
5461
import java.io.OutputStream;
5562
import java.util.ArrayList;
5663
import java.util.Arrays;
@@ -93,6 +100,7 @@ public SearchHits(SearchHit[] hits, @Nullable TotalHits totalHits, float maxScor
93100
this(hits, totalHits, maxScore, null, null, null);
94101
}
95102

103+
@SuppressForbidden(reason = "serialization of object to protobuf")
96104
public SearchHits(
97105
SearchHit[] hits,
98106
@Nullable TotalHits totalHits,
@@ -109,18 +117,18 @@ public SearchHits(
109117
this.collapseValues = collapseValues;
110118
if (FeatureFlags.isEnabled(FeatureFlags.PROTOBUF_SETTING)) {
111119
List<FetchSearchResultProto.SearchHit> searchHitList = new ArrayList<>();
112-
for (int i = 0; i < hits.length; i++) {
120+
for (SearchHit hit : hits) {
113121
FetchSearchResultProto.SearchHit.Builder searchHitBuilder = FetchSearchResultProto.SearchHit.newBuilder();
114-
if (hits[i].getIndex() != null) {
115-
searchHitBuilder.setIndex(hits[i].getIndex());
122+
if (hit.getIndex() != null) {
123+
searchHitBuilder.setIndex(hit.getIndex());
116124
}
117-
searchHitBuilder.setId(hits[i].getId());
118-
searchHitBuilder.setScore(hits[i].getScore());
119-
searchHitBuilder.setSeqNo(hits[i].getSeqNo());
120-
searchHitBuilder.setPrimaryTerm(hits[i].getPrimaryTerm());
121-
searchHitBuilder.setVersion(hits[i].getVersion());
122-
if (hits[i].getSourceRef() != null) {
123-
searchHitBuilder.setSource(ByteString.copyFrom(hits[i].getSourceRef().toBytesRef().bytes));
125+
searchHitBuilder.setId(hit.getId());
126+
searchHitBuilder.setScore(hit.getScore());
127+
searchHitBuilder.setSeqNo(hit.getSeqNo());
128+
searchHitBuilder.setPrimaryTerm(hit.getPrimaryTerm());
129+
searchHitBuilder.setVersion(hit.getVersion());
130+
if (hit.getSourceRef() != null) {
131+
searchHitBuilder.setSource(ByteString.copyFrom(hit.getSourceRef().toBytesRef().bytes));
124132
}
125133
searchHitList.add(searchHitBuilder.build());
126134
}
@@ -135,8 +143,25 @@ public SearchHits(
135143
searchHitsBuilder.setMaxScore(maxScore);
136144
searchHitsBuilder.addAllHits(searchHitList);
137145
searchHitsBuilder.setTotalHits(totalHitsBuilder.build());
146+
if (sortFields != null && sortFields.length > 0) {
147+
for (SortField sortField : sortFields) {
148+
FetchSearchResultProto.SortField.Builder sortFieldBuilder = FetchSearchResultProto.SortField.newBuilder();
149+
sortFieldBuilder.setField(sortField.getField());
150+
sortFieldBuilder.setType(FetchSearchResultProto.SortField.Type.valueOf(sortField.getType().name()));
151+
searchHitsBuilder.addSortFields(sortFieldBuilder.build());
152+
}
153+
}
138154
if (collapseField != null) {
139155
searchHitsBuilder.setCollapseField(collapseField);
156+
for (Object value : collapseValues) {
157+
ByteArrayOutputStream bos = new ByteArrayOutputStream();
158+
try (ObjectOutputStream stream = new ObjectOutputStream(bos)) {
159+
stream.writeObject(value);
160+
searchHitsBuilder.addCollapseValues(ByteString.copyFrom(bos.toByteArray()));
161+
} catch (IOException e) {
162+
throw new OpenSearchException(e);
163+
}
164+
}
140165
}
141166
this.searchHitsProto = searchHitsBuilder.build();
142167
}
@@ -164,6 +189,7 @@ public SearchHits(StreamInput in) throws IOException {
164189
collapseValues = in.readOptionalArray(Lucene::readSortValue, Object[]::new);
165190
}
166191

192+
@SuppressForbidden(reason = "serialization of object to protobuf")
167193
public SearchHits(byte[] in) throws IOException {
168194
this.searchHitsProto = org.opensearch.server.proto.FetchSearchResultProto.SearchHits.parseFrom(in);
169195
this.hits = new SearchHit[this.searchHitsProto.getHitsCount()];
@@ -175,10 +201,21 @@ public SearchHits(byte[] in) throws IOException {
175201
Relation.valueOf(this.searchHitsProto.getTotalHits().getRelation().toString())
176202
);
177203
this.maxScore = this.searchHitsProto.getMaxScore();
204+
this.sortFields = this.searchHitsProto.getSortFieldsList()
205+
.stream()
206+
.map(sortField -> new SortField(sortField.getField(), SortField.Type.valueOf(sortField.getType().toString())))
207+
.toArray(SortField[]::new);
178208
this.collapseField = this.searchHitsProto.getCollapseField();
179-
// Below fields are set to null currently, support to be added in the future
180-
this.collapseValues = null;
181-
this.sortFields = null;
209+
this.collapseValues = new Object[this.searchHitsProto.getCollapseValuesCount()];
210+
for (int i = 0; i < this.searchHitsProto.getCollapseValuesCount(); i++) {
211+
ByteString collapseValue = this.searchHitsProto.getCollapseValues(i);
212+
InputStream is = new ByteArrayInputStream(collapseValue.toByteArray());
213+
try (ObjectInputStream ois = new ObjectInputStream(is)) {
214+
this.collapseValues[i] = ois.readObject();
215+
} catch (ClassNotFoundException e) {
216+
throw new OpenSearchException(e);
217+
}
218+
}
182219
}
183220

184221
@Override

0 commit comments

Comments
 (0)