Introduce protobuf serialization for objects - FetchSearchResult, SearchHits, SearchHit

buildSrc/src/main/groovy/org/opensearch/gradle/precommit/LicenseHeadersTask.groovy

@@ -149,6 +149,8 @@ class LicenseHeadersTask extends AntTask {
                              licenseFamilyName:     "Generated") {
                // parsers generated by antlr
                pattern(substring: "ANTLR GENERATED CODE")
+               // Protobuf
+               pattern(substring: "Generated by the protocol buffer compiler")
             }
 
             // Vendored Code

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

@@ -98,46 +98,74 @@
  * @opensearch.api
  */
 @PublicApi(since = "1.0.0")
-public final class SearchHit implements Writeable, ToXContentObject, Iterable<DocumentField> {
+public class SearchHit implements Writeable, ToXContentObject, Iterable<DocumentField> {
 
-    private final transient int docId;
+    protected transient int docId;
 
     private static final float DEFAULT_SCORE = Float.NaN;
-    private float score = DEFAULT_SCORE;
+    protected float score = DEFAULT_SCORE;
 
-    private final Text id;
+    @Nullable
+    protected Text id;
 
-    private final NestedIdentity nestedIdentity;
+    @Nullable
+    protected NestedIdentity nestedIdentity;
 
-    private long version = -1;
-    private long seqNo = SequenceNumbers.UNASSIGNED_SEQ_NO;
-    private long primaryTerm = SequenceNumbers.UNASSIGNED_PRIMARY_TERM;
+    protected long version = -1;
+    protected long seqNo = SequenceNumbers.UNASSIGNED_SEQ_NO;
+    protected long primaryTerm = SequenceNumbers.UNASSIGNED_PRIMARY_TERM;
 
-    private BytesReference source;
+    @Nullable
+    protected BytesReference source;
 
-    private Map<String, DocumentField> documentFields;
-    private final Map<String, DocumentField> metaFields;
+    protected Map<String, DocumentField> documentFields;
+    protected Map<String, DocumentField> metaFields;
 
-    private Map<String, HighlightField> highlightFields = null;
+    @Nullable
+    protected Map<String, HighlightField> highlightFields = null;
 
-    private SearchSortValues sortValues = SearchSortValues.EMPTY;
+    protected SearchSortValues sortValues = SearchSortValues.EMPTY;
 
-    private Map<String, Float> matchedQueries = new HashMap<>();
+    protected Map<String, Float> matchedQueries = new HashMap<>();
 
-    private Explanation explanation;
+    @Nullable
+    protected Explanation explanation;
 
     @Nullable
-    private SearchShardTarget shard;
+    protected SearchShardTarget shard;
 
     // These two fields normally get set when setting the shard target, so they hold the same values as the target thus don't get
     // serialized over the wire. When parsing hits back from xcontent though, in most of the cases (whenever explanation is disabled)
     // we can't rebuild the shard target object so we need to set these manually for users retrieval.
-    private transient String index;
-    private transient String clusterAlias;
+    protected transient String index;
+    protected transient String clusterAlias;
 
     private Map<String, Object> sourceAsMap;
 
-    private Map<String, SearchHits> innerHits;
+    @Nullable
+    protected Map<String, SearchHits> innerHits;
+
+    public SearchHit(SearchHit hit) {
+        this.docId = hit.docId;
+        this.id = hit.id;
+        this.nestedIdentity = hit.nestedIdentity;
+        this.version = hit.version;
+        this.seqNo = hit.seqNo;
+        this.primaryTerm = hit.primaryTerm;
+        this.source = hit.source;
+        this.documentFields = hit.documentFields;
+        this.metaFields = hit.metaFields;
+        this.highlightFields = hit.highlightFields;
+        this.sortValues = hit.sortValues;
+        this.matchedQueries = hit.matchedQueries;
+        this.explanation = hit.explanation;
+        this.shard = hit.shard;
+        this.index = hit.index;
+        this.clusterAlias = hit.clusterAlias;
+        this.innerHits = hit.innerHits;
+        this.score = hit.score;
+        this.sourceAsMap = hit.sourceAsMap;
+    }
 
     // used only in tests
     public SearchHit(int docId) {
@@ -236,7 +264,9 @@ public SearchHit(StreamInput in) throws IOException {
         }
     }
 
-    private static final Text SINGLE_MAPPING_TYPE = new Text(MapperService.SINGLE_MAPPING_NAME);
+    protected SearchHit() {}
+
+    protected static final Text SINGLE_MAPPING_TYPE = new Text(MapperService.SINGLE_MAPPING_NAME);
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
@@ -993,7 +1023,7 @@ private void buildExplanation(XContentBuilder builder, Explanation explanation)
 
     @Override
     public boolean equals(Object obj) {
-        if (obj == null || getClass() != obj.getClass()) {
+        if (!(obj instanceof SearchHit)) {
             return false;
         }
         SearchHit other = (SearchHit) obj;
@@ -1057,7 +1087,7 @@ public NestedIdentity(String field, int offset, NestedIdentity child) {
             this.child = child;
         }
 
-        NestedIdentity(StreamInput in) throws IOException {
+        public NestedIdentity(StreamInput in) throws IOException {
             field = in.readOptionalText();
             offset = in.readInt();
             child = in.readOptionalWriteable(NestedIdentity::new);

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

@@ -61,7 +61,7 @@
  * @opensearch.api
  */
 @PublicApi(since = "1.0.0")
-public final class SearchHits implements Writeable, ToXContentFragment, Iterable<SearchHit> {
+public class SearchHits implements Writeable, ToXContentFragment, Iterable<SearchHit> {
     public static SearchHits empty() {
         return empty(true);
     }
@@ -72,15 +72,25 @@ public static SearchHits empty(boolean withTotalHits) {
 
     public static final SearchHit[] EMPTY = new SearchHit[0];
 
-    private final SearchHit[] hits;
-    private final TotalHits totalHits;
-    private final float maxScore;
+    protected SearchHit[] hits;
+    protected float maxScore;
     @Nullable
-    private final SortField[] sortFields;
+    protected TotalHits totalHits;
     @Nullable
-    private final String collapseField;
+    protected SortField[] sortFields;
     @Nullable
-    private final Object[] collapseValues;
+    protected String collapseField;
+    @Nullable
+    protected Object[] collapseValues;
+
+    public SearchHits(SearchHits sHits) {
+        this.hits = sHits.hits;
+        this.totalHits = sHits.totalHits;
+        this.maxScore = sHits.maxScore;
+        this.sortFields = sHits.sortFields;
+        this.collapseField = sHits.collapseField;
+        this.collapseValues = sHits.collapseValues;
+    }
 
     public SearchHits(SearchHit[] hits, @Nullable TotalHits totalHits, float maxScore) {
         this(hits, totalHits, maxScore, null, null, null);
@@ -124,6 +134,8 @@ public SearchHits(StreamInput in) throws IOException {
         collapseValues = in.readOptionalArray(Lucene::readSortValue, Object[]::new);
     }
 
+    protected SearchHits() {}
+
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         final boolean hasTotalHits = totalHits != null;
@@ -288,7 +300,7 @@ public static SearchHits fromXContent(XContentParser parser) throws IOException
 
     @Override
     public boolean equals(Object obj) {
-        if (obj == null || getClass() != obj.getClass()) {
+        if (!(obj instanceof SearchHits)) {
             return false;
         }
         SearchHits other = (SearchHits) obj;

server/src/main/java/org/opensearch/search/SearchShardTarget.java

@@ -58,6 +58,8 @@ public final class SearchShardTarget implements Writeable, Comparable<SearchShar
     // original indices are only needed in the coordinating node throughout the search request execution.
     // no need to serialize them as part of SearchShardTarget.
     private final transient OriginalIndices originalIndices;
+
+    @Nullable
     private final String clusterAlias;
 
     public SearchShardTarget(StreamInput in) throws IOException {
@@ -71,6 +73,13 @@ public SearchShardTarget(StreamInput in) throws IOException {
         clusterAlias = in.readOptionalString();
     }
 
+    public SearchShardTarget(String nodeId, ShardId shardId, String clusterAlias) {
+        this.nodeId = new Text(nodeId);
+        this.shardId = shardId;
+        this.originalIndices = null;
+        this.clusterAlias = clusterAlias;
+    }
+
     public SearchShardTarget(String nodeId, ShardId shardId, @Nullable String clusterAlias, OriginalIndices originalIndices) {
         this.nodeId = nodeId == null ? null : new Text(nodeId);
         this.shardId = shardId;

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

@@ -67,6 +67,11 @@ public class SearchSortValues implements ToXContentFragment, Writeable {
         this.rawSortValues = EMPTY_ARRAY;
     }
 
+    public SearchSortValues(Object[] formattedSortValues, Object[] rawSortValues) {
+        this.formattedSortValues = formattedSortValues;
+        this.rawSortValues = rawSortValues;
+    }
+
     public SearchSortValues(Object[] rawSortValues, DocValueFormat[] sortValueFormats) {
         Objects.requireNonNull(rawSortValues);
         Objects.requireNonNull(sortValueFormats);
@@ -89,7 +94,7 @@ public SearchSortValues(Object[] rawSortValues, DocValueFormat[] sortValueFormat
         }
     }
 
-    SearchSortValues(StreamInput in) throws IOException {
+    public SearchSortValues(StreamInput in) throws IOException {
         this.formattedSortValues = in.readArray(Lucene::readSortValue, Object[]::new);
         this.rawSortValues = in.readArray(Lucene::readSortValue, Object[]::new);
     }

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

@@ -50,11 +50,11 @@
  * @opensearch.api
  */
 @PublicApi(since = "1.0.0")
-public final class FetchSearchResult extends SearchPhaseResult {
+public class FetchSearchResult extends SearchPhaseResult {
 
-    private SearchHits hits;
+    protected SearchHits hits;
     // client side counter
-    private transient int counter;
+    protected transient int counter;
 
     public FetchSearchResult() {}
 

server/src/main/java/org/opensearch/search/fetch/subphase/highlight/HighlightField.java

@@ -32,6 +32,7 @@
 
 package org.opensearch.search.fetch.subphase.highlight;
 
+import org.opensearch.common.Nullable;
 import org.opensearch.common.annotation.PublicApi;
 import org.opensearch.core.common.ParsingException;
 import org.opensearch.core.common.io.stream.StreamInput;
@@ -60,6 +61,7 @@ public class HighlightField implements ToXContentFragment, Writeable {
 
     private String name;
 
+    @Nullable
     private Text[] fragments;
 
     public HighlightField(StreamInput in) throws IOException {

server/src/main/java/org/opensearch/transport/TransportSerializationException.java

@@ -47,6 +47,10 @@ public TransportSerializationException(StreamInput in) throws IOException {
         super(in);
     }
 
+    public TransportSerializationException(String msg) {
+        super(msg);
+    }
+
     public TransportSerializationException(String msg, Throwable cause) {
         super(msg, cause);
     }

server/src/main/java/org/opensearch/transport/protobuf/FetchSearchResultProtobuf.java

@@ -0,0 +1,51 @@
+/*
+ * 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.protobuf;
+
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.core.common.io.stream.StreamOutput;
+import org.opensearch.search.fetch.FetchSearchResult;
+import org.opensearch.proto.search.fetch.FetchSearchResultProtoDef.FetchSearchResultProto;
+
+import java.io.IOException;
+
+/**
+ * FetchSearchResult which leverages protobuf for transport layer serialization.
+ * @opensearch.internal
+ */
+public class FetchSearchResultProtobuf extends FetchSearchResult {
+    public FetchSearchResultProtobuf(StreamInput in) throws IOException {
+        fromProtobufStream(in);
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        toProtobufStream(out);
+    }
+
+    public void toProtobufStream(StreamOutput out) throws IOException {
+        toProto().writeTo(out);
+    }
+
+    public void fromProtobufStream(StreamInput in) throws IOException {
+        FetchSearchResultProto proto = FetchSearchResultProto.parseFrom(in);
+        fromProto(proto);
+    }
+
+    FetchSearchResultProto toProto() {
+        FetchSearchResultProto.Builder builder = FetchSearchResultProto.newBuilder()
+            .setHits(new SearchHitsProtobuf(hits).toProto())
+            .setCounter(this.counter);
+        return builder.build();
+    }
+
+    void fromProto(FetchSearchResultProto proto) {
+        hits = new SearchHitsProtobuf(proto.getHits());
+    }
+}