Converting SearchHits to a proto message

Signed-off-by: Vacha Shah <[email protected]>
VachaShah · Dec 12, 2023 · 9ec8a1a · 9ec8a1a
1 parent 069828f
commit 9ec8a1a
Show file tree

Hide file tree

Showing 6 changed files with 76 additions and 26 deletions.
diff --git a/server/src/main/java/org/opensearch/search/fetch/FetchSearchResult.java b/server/src/main/java/org/opensearch/search/fetch/FetchSearchResult.java
@@ -55,7 +55,6 @@
  */
 public final class FetchSearchResult extends SearchPhaseResult {
 
-    // TODO: Write SearchHits as a proto message
     private SearchHits hits;
     // client side counter
     private transient int counter;

diff --git a/server/src/main/java/org/opensearch/search/internal/ProtobufShardSearchRequest.java b/server/src/main/java/org/opensearch/search/internal/ProtobufShardSearchRequest.java
@@ -234,7 +234,7 @@ private ProtobufShardSearchRequest(
 
         if (scroll != null) {
             // TODO: Write Scroll as a proto message
-            builder.setScroll(ByteString.copyFrom(convertToBytes(scroll)));
+            builder.setScroll(ShardSearchRequestProto.Scroll.newBuilder().setKeepAlive(scroll.keepAlive().getStringRep()));
         }
         builder.setNowInMillis(nowInMillis);
 
@@ -377,13 +377,6 @@ public boolean allowPartialSearchResults() {
     }
 
     public Scroll scroll() {
-        ByteArrayInputStream in = new ByteArrayInputStream(this.shardSearchRequestProto.getScroll().toByteArray());
-        try (ObjectInputStream is = new ObjectInputStream(in)) {
-            return (Scroll) is.readObject();
-        } catch (ClassNotFoundException | IOException e) {
-            // TODO Auto-generated catch block
-            e.printStackTrace();
-        }
         return null;
     }
 

diff --git a/server/src/main/java/org/opensearch/search/query/QuerySearchResult.java b/server/src/main/java/org/opensearch/search/query/QuerySearchResult.java
@@ -138,7 +138,7 @@ public QuerySearchResult(ShardSearchContextId contextId, SearchShardTarget shard
                         .setIndexName(shardTarget.getShardId().getIndexName())
                         .setIndexUUID(shardTarget.getShardId().getIndex().getUUID())
                         .build();
-        QuerySearchResultProto.QuerySearchResult.SearchShardTarget searchShardTarget = QuerySearchResultProto.QuerySearchResult.SearchShardTarget.newBuilder()
+        QuerySearchResultProto.SearchShardTarget searchShardTarget = QuerySearchResultProto.SearchShardTarget.newBuilder()
                                                                                             .setNodeId(shardTarget.getNodeId())
                                                                                             .setShardId(shardIdProto)
                                                                                             .setClusterAlias(shardTarget.getClusterAlias())

diff --git a/server/src/main/proto/server/FetchSearchResultProto.proto b/server/src/main/proto/server/FetchSearchResultProto.proto
@@ -13,10 +13,63 @@ syntax = "proto3";
 package org.opensearch.server.proto;
 
 import "server/ShardSearchRequestProto.proto";
+import "server/QuerySearchResultProto.proto";
 
 option java_outer_classname = "FetchSearchResultProto";
 
 message FetchSearchResult {
   ShardSearchContextId contextId = 1;
-  optional bytes hits = 2;
+  optional SearchHits hits = 2;
+}
+
+message SearchHits {
+  TotalHits totalHits = 1;
+  float maxScore = 2;
+  int32 size = 3;
+  repeated SearchHit hits = 4;
+  /* optional repeated SortField sortFields = 5; */
+  optional string collapseField = 5;
+  repeated bytes collapseValues = 6;
+}
+
+message SearchHit {
+  int32 docId = 1;
+  float score = 2;
+  string id = 3;
+  NestedIdentity nestedIdentity = 4;
+  int64 version = 5;
+  int64 seqNo = 6;
+  int64 primaryTerm = 7;
+  bytes source = 8;
+  map<string, DocumentField> documentFields = 9;
+  map<string, DocumentField> metaFields = 10;
+  map<string, HighlightField> highlightFields = 11;
+  SearchSortValues sortValues = 12;
+  repeated string matchedQueries = 13;
+  /* Explanation explanation = 14;*/
+  SearchShardTarget shard = 14;
+  string index = 15;
+  string clusterAlias = 16;
+  map<string, bytes> sourceAsMap = 17;
+
+  message NestedIdentity {
+    string field = 1;
+    int32 offset = 2;
+    NestedIdentity child = 3;
+  }
+
+  message DocumentField {
+    string name = 1;
+    repeated bytes values = 2;
+  }
+
+  message HighlightField {
+    string name = 1;
+    repeated string fragments = 2;
+  }
+
+  message SearchSortValues {
+    repeated bytes formattedSortValues = 1;
+    repeated bytes rawSortValues = 2;
+  }
 }
diff --git a/server/src/main/proto/server/QuerySearchResultProto.proto b/server/src/main/proto/server/QuerySearchResultProto.proto
@@ -41,16 +41,6 @@ message QuerySearchResult {
     float maxScore = 2;
   }
 
-  message TotalHits {
-    int64 value = 1;
-    Relation relation = 2;
-
-    enum Relation {
-        EQUAL_TO = 0;
-        GREATER_THAN_OR_EQUAL_TO = 1;
-    }
-  }
-
   message TopDocs {
     TotalHits totalHits = 1;
     repeated ScoreDoc scoreDocs = 2;
@@ -70,9 +60,20 @@ message QuerySearchResult {
     }
   }
 
-  message SearchShardTarget {
-    string nodeId = 1;
-    ShardId shardId = 2;
-    string clusterAlias = 3;
+}
+
+message SearchShardTarget {
+  string nodeId = 1;
+  ShardId shardId = 2;
+  string clusterAlias = 3;
+}
+
+message TotalHits {
+  int64 value = 1;
+  Relation relation = 2;
+
+  enum Relation {
+      EQUAL_TO = 0;
+      GREATER_THAN_OR_EQUAL_TO = 1;
   }
 }
diff --git a/server/src/main/proto/server/ShardSearchRequestProto.proto b/server/src/main/proto/server/ShardSearchRequestProto.proto
@@ -26,7 +26,7 @@ message ShardSearchRequest {
   bool allowPartialSearchResults = 9;
   repeated string indexRoutings = 10;
   string preference = 11;
-  bytes scroll = 12;
+  Scroll scroll = 12;
   int64 nowInMillis = 13;
   optional string clusterAlias = 14;
   optional ShardSearchContextId readerId = 15;
@@ -53,6 +53,10 @@ message ShardId {
   string indexUUID = 4;
 }
 
+message Scroll {
+  string keepAlive = 1; 
+}
+
 message OriginalIndices {
   repeated string indices = 1;
   IndicesOptions indicesOptions = 2;