Adding protos for search classes

CHANGELOG.md

@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ### Added
 - Add useCompoundFile index setting ([#13478](https://github.com/opensearch-project/OpenSearch/pull/13478))
 - Make outbound side of transport protocol dependent ([#13293](https://github.com/opensearch-project/OpenSearch/pull/13293))
+- Add proto structures for Query and Fetch search result. ([#13178](https://github.com/opensearch-project/OpenSearch/pull/13178))
 
 ### Dependencies
 - Bump `com.github.spullara.mustache.java:compiler` from 0.9.10 to 0.9.13 ([#13329](https://github.com/opensearch-project/OpenSearch/pull/13329), [#13559](https://github.com/opensearch-project/OpenSearch/pull/13559))

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

@@ -238,7 +238,15 @@ private boolean inspectable(ExecutableElement executable) {
      */
     private boolean inspectable(Element element) {
         final PackageElement pckg = processingEnv.getElementUtils().getPackageOf(element);
-        return pckg.getQualifiedName().toString().startsWith(OPENSEARCH_PACKAGE);
+        return pckg.getQualifiedName().toString().startsWith(OPENSEARCH_PACKAGE)
+            && !element.getEnclosingElement()
+                .getAnnotationMirrors()
+                .stream()
+                .anyMatch(
+                    m -> m.getAnnotationType()
+                        .toString() /* ClassSymbol.toString() returns class name */
+                        .equalsIgnoreCase("javax.annotation.Generated")
+                );
     }
 
     /**

server/build.gradle

@@ -380,6 +380,7 @@ tasks.named("licenseHeaders").configure {
     excludes << 'org/opensearch/client/documentation/placeholder.txt'
     // Ignore for protobuf generated code
     excludes << 'org/opensearch/extensions/proto/*'
+    excludes << 'org/opensearch/server/proto/*'
 }
 
 tasks.test {

server/src/main/java/org/opensearch/common/document/serializer/DocumentFieldDeserializer.java

@@ -0,0 +1,22 @@
+/*
+ * 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.common.document.serializer;
+
+import org.opensearch.common.document.DocumentField;
+
+import java.io.IOException;
+
+/**
+ * Deserializer for {@link DocumentField} which can be implemented for different types of serde mechanisms.
+ */
+public interface DocumentFieldDeserializer<T> {
+
+    DocumentField createDocumentField(T inputStream) throws IOException;
+
+}

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

@@ -0,0 +1,10 @@
+/*
+ * 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.
+ */
+
+/** Serializer package for documents. */
+package org.opensearch.common.document.serializer;

server/src/main/java/org/opensearch/common/document/serializer/protobuf/DocumentFieldProtobufDeserializer.java

@@ -0,0 +1,144 @@
+/*
+ * 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.common.document.serializer.protobuf;
+
+import com.google.protobuf.ByteString;
+import org.opensearch.OpenSearchException;
+import org.opensearch.common.document.DocumentField;
+import org.opensearch.common.document.serializer.DocumentFieldDeserializer;
+import org.opensearch.core.common.text.Text;
+import org.opensearch.server.proto.FetchSearchResultProto;
+import org.opensearch.server.proto.FetchSearchResultProto.DocumentFieldValue;
+import org.opensearch.server.proto.FetchSearchResultProto.DocumentFieldValue.Builder;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.time.Instant;
+import java.time.ZoneId;
+import java.time.ZonedDateTime;
+import java.util.ArrayList;
+import java.util.Date;
+import java.util.HashMap;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Deserializer for {@link DocumentField} to/from protobuf.
+ */
+public class DocumentFieldProtobufDeserializer implements DocumentFieldDeserializer<InputStream> {
+
+    private FetchSearchResultProto.SearchHit.DocumentField documentField;
+
+    @Override
+    public DocumentField createDocumentField(InputStream inputStream) throws IOException {
+        documentField = FetchSearchResultProto.SearchHit.DocumentField.parseFrom(inputStream);
+        String name = documentField.getName();
+        List<Object> values = new ArrayList<>();
+        for (FetchSearchResultProto.DocumentFieldValue value : documentField.getValuesList()) {
+            values.add(readDocumentFieldValueFromProtobuf(value));
+        }
+        return new DocumentField(name, values);
+    }
+
+    private Object readDocumentFieldValueFromProtobuf(FetchSearchResultProto.DocumentFieldValue documentFieldValue) throws IOException {
+        if (documentFieldValue.hasValueString()) {
+            return documentFieldValue.getValueString();
+        } else if (documentFieldValue.hasValueInt()) {
+            return documentFieldValue.getValueInt();
+        } else if (documentFieldValue.hasValueLong()) {
+            return documentFieldValue.getValueLong();
+        } else if (documentFieldValue.hasValueFloat()) {
+            return documentFieldValue.getValueFloat();
+        } else if (documentFieldValue.hasValueDouble()) {
+            return documentFieldValue.getValueDouble();
+        } else if (documentFieldValue.hasValueBool()) {
+            return documentFieldValue.getValueBool();
+        } else if (documentFieldValue.getValueByteArrayList().size() > 0) {
+            return documentFieldValue.getValueByteArrayList().toArray();
+        } else if (documentFieldValue.getValueArrayListList().size() > 0) {
+            List<Object> list = new ArrayList<>();
+            for (FetchSearchResultProto.DocumentFieldValue value : documentFieldValue.getValueArrayListList()) {
+                list.add(readDocumentFieldValueFromProtobuf(value));
+            }
+            return list;
+        } else if (documentFieldValue.getValueMapMap().size() > 0) {
+            Map<String, Object> map = Map.of();
+            for (Map.Entry<String, FetchSearchResultProto.DocumentFieldValue> entrySet : documentFieldValue.getValueMapMap().entrySet()) {
+                map.put(entrySet.getKey(), readDocumentFieldValueFromProtobuf(entrySet.getValue()));
+            }
+            return map;
+        } else if (documentFieldValue.hasValueDate()) {
+            return new Date(documentFieldValue.getValueDate());
+        } else if (documentFieldValue.hasValueZonedDate() && documentFieldValue.hasValueZonedTime()) {
+            return ZonedDateTime.ofInstant(
+                Instant.ofEpochMilli(documentFieldValue.getValueZonedTime()),
+                ZoneId.of(documentFieldValue.getValueZonedDate())
+            );
+        } else if (documentFieldValue.hasValueText()) {
+            return new Text(documentFieldValue.getValueText());
+        } else {
+            throw new IOException("Can't read generic value of type [" + documentFieldValue + "]");
+        }
+    }
+
+    public static DocumentFieldValue.Builder convertDocumentFieldValueToProto(Object value, Builder valueBuilder) {
+        if (value == null) {
+            // null is not allowed in protobuf, so we use a special string to represent null
+            return valueBuilder.setValueString("null");
+        }
+        Class type = value.getClass();
+        if (type == String.class) {
+            valueBuilder.setValueString((String) value);
+        } else if (type == Integer.class) {
+            valueBuilder.setValueInt((Integer) value);
+        } else if (type == Long.class) {
+            valueBuilder.setValueLong((Long) value);
+        } else if (type == Float.class) {
+            valueBuilder.setValueFloat((Float) value);
+        } else if (type == Double.class) {
+            valueBuilder.setValueDouble((Double) value);
+        } else if (type == Boolean.class) {
+            valueBuilder.setValueBool((Boolean) value);
+        } else if (type == byte[].class) {
+            valueBuilder.addValueByteArray(ByteString.copyFrom((byte[]) value));
+        } else if (type == List.class) {
+            List<Object> list = (List<Object>) value;
+            for (Object listValue : list) {
+                valueBuilder.addValueArrayList(convertDocumentFieldValueToProto(listValue, valueBuilder));
+            }
+        } else if (type == Map.class || type == HashMap.class || type == LinkedHashMap.class) {
+            Map<String, Object> map = (Map<String, Object>) value;
+            for (Map.Entry<String, Object> entry : map.entrySet()) {
+                valueBuilder.putValueMap(entry.getKey(), convertDocumentFieldValueToProto(entry.getValue(), valueBuilder).build());
+            }
+        } else if (type == Date.class) {
+            valueBuilder.setValueDate(((Date) value).getTime());
+        } else if (type == ZonedDateTime.class) {
+            valueBuilder.setValueZonedDate(((ZonedDateTime) value).getZone().getId());
+            valueBuilder.setValueZonedTime(((ZonedDateTime) value).toInstant().toEpochMilli());
+        } else if (type == Text.class) {
+            valueBuilder.setValueText(((Text) value).string());
+        } else {
+            throw new OpenSearchException("Can't convert generic value of type [" + type + "] to protobuf");
+        }
+        return valueBuilder;
+    }
+
+    public static FetchSearchResultProto.SearchHit.DocumentField convertDocumentFieldToProto(DocumentField documentField) {
+        FetchSearchResultProto.SearchHit.DocumentField.Builder builder = FetchSearchResultProto.SearchHit.DocumentField.newBuilder();
+        builder.setName(documentField.getName());
+        for (Object value : documentField.getValues()) {
+            FetchSearchResultProto.DocumentFieldValue.Builder valueBuilder = FetchSearchResultProto.DocumentFieldValue.newBuilder();
+            builder.addValues(convertDocumentFieldValueToProto(value, valueBuilder));
+        }
+        return builder.build();
+    }
+
+}

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

@@ -0,0 +1,10 @@
+/*
+ * 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.
+ */
+
+/** Protobuf Serializer package for documents. */
+package org.opensearch.common.document.serializer.protobuf;

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

@@ -67,6 +67,11 @@ public class FeatureFlags {
      */
     public static final String PLUGGABLE_CACHE = "opensearch.experimental.feature.pluggable.caching.enabled";
 
+    /**
+     * Gates the functionality of integrating protobuf within search API and node-to-node communication.
+     */
+    public static final String PROTOBUF = "opensearch.experimental.feature.search_with_protobuf.enabled";
+
     public static final Setting<Boolean> REMOTE_STORE_MIGRATION_EXPERIMENTAL_SETTING = Setting.boolSetting(
         REMOTE_STORE_MIGRATION_EXPERIMENTAL,
         false,
@@ -93,14 +98,17 @@ public class FeatureFlags {
 
     public static final Setting<Boolean> PLUGGABLE_CACHE_SETTING = Setting.boolSetting(PLUGGABLE_CACHE, false, Property.NodeScope);
 
+    public static final Setting<Boolean> PROTOBUF_SETTING = Setting.boolSetting(PROTOBUF, false, Property.NodeScope, Property.Dynamic);
+
     private static final List<Setting<Boolean>> ALL_FEATURE_FLAG_SETTINGS = List.of(
         REMOTE_STORE_MIGRATION_EXPERIMENTAL_SETTING,
         EXTENSIONS_SETTING,
         IDENTITY_SETTING,
         TELEMETRY_SETTING,
         DATETIME_FORMATTER_CACHING_SETTING,
         WRITEABLE_REMOTE_INDEX_SETTING,
-        PLUGGABLE_CACHE_SETTING
+        PLUGGABLE_CACHE_SETTING,
+        PROTOBUF_SETTING
     );
     /**
      * Should store the settings from opensearch.yml.

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

@@ -67,6 +67,16 @@
         this.rawSortValues = EMPTY_ARRAY;
     }
 
+    public SearchSortValues(Object[] sortValues, Object[] rawSortValues) {
+        Objects.requireNonNull(rawSortValues);
+        Objects.requireNonNull(sortValues);
+        if (rawSortValues.length != sortValues.length) {
+            throw new IllegalArgumentException("formattedSortValues and sortValues must hold the same number of items");
+        }
+        this.formattedSortValues = sortValues;
+        this.rawSortValues = rawSortValues;
+    }
+
     public SearchSortValues(Object[] rawSortValues, DocValueFormat[] sortValueFormats) {
         Objects.requireNonNull(rawSortValues);
         Objects.requireNonNull(sortValueFormats);

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

@@ -33,6 +33,7 @@
 package org.opensearch.search.fetch;
 
 import org.opensearch.common.annotation.PublicApi;
+import org.opensearch.common.util.FeatureFlags;
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.core.common.io.stream.StreamOutput;
 import org.opensearch.search.SearchHit;
@@ -41,8 +42,13 @@
 import org.opensearch.search.SearchShardTarget;
 import org.opensearch.search.internal.ShardSearchContextId;
 import org.opensearch.search.query.QuerySearchResult;
+import org.opensearch.search.serializer.protobuf.SearchHitsProtobufDeserializer;
+import org.opensearch.server.proto.FetchSearchResultProto;
+import org.opensearch.server.proto.ShardSearchRequestProto;
 
+import java.io.ByteArrayInputStream;
 import java.io.IOException;
+import java.io.InputStream;
 
 /**
  * Result from a fetch
@@ -56,6 +62,8 @@
     // client side counter
     private transient int counter;
 
+    private FetchSearchResultProto.FetchSearchResult fetchSearchResultProto;
+
     public FetchSearchResult() {}
 
     public FetchSearchResult(StreamInput in) throws IOException {
@@ -64,9 +72,24 @@
         hits = new SearchHits(in);
     }
 
+    public FetchSearchResult(InputStream in) throws IOException {
+        this.fetchSearchResultProto = FetchSearchResultProto.FetchSearchResult.parseFrom(in);
+        contextId = new ShardSearchContextId(
+            this.fetchSearchResultProto.getContextId().getSessionId(),
+            this.fetchSearchResultProto.getContextId().getId()
+        );
+        SearchHitsProtobufDeserializer protobufSerializer = new SearchHitsProtobufDeserializer();
+        hits = protobufSerializer.createSearchHits(new ByteArrayInputStream(this.fetchSearchResultProto.getHits().toByteArray()));
+    }
+
     public FetchSearchResult(ShardSearchContextId id, SearchShardTarget shardTarget) {
         this.contextId = id;
         setSearchShardTarget(shardTarget);
+        this.fetchSearchResultProto = FetchSearchResultProto.FetchSearchResult.newBuilder()
+            .setContextId(
+                ShardSearchRequestProto.ShardSearchContextId.newBuilder().setSessionId(id.getSessionId()).setId(id.getId()).build()
+            )
+            .build();
     }
 
     @Override
@@ -82,6 +105,11 @@
     public void hits(SearchHits hits) {
         assert assertNoSearchTarget(hits);
         this.hits = hits;
+        if (FeatureFlags.isEnabled(FeatureFlags.PROTOBUF_SETTING) && this.fetchSearchResultProto != null) {
+            this.fetchSearchResultProto = this.fetchSearchResultProto.toBuilder()
+                .setHits(SearchHitsProtobufDeserializer.convertHitsToProto(hits))
+                .build();
+        }
     }
 
     private boolean assertNoSearchTarget(SearchHits hits) {
@@ -92,6 +120,16 @@
     }
 
     public SearchHits hits() {
+        if (FeatureFlags.isEnabled(FeatureFlags.PROTOBUF_SETTING) && this.fetchSearchResultProto != null) {
+            SearchHits hits;
+            try {
+                SearchHitsProtobufDeserializer protobufSerializer = new SearchHitsProtobufDeserializer();
+                hits = protobufSerializer.createSearchHits(new ByteArrayInputStream(this.fetchSearchResultProto.getHits().toByteArray()));
+                return hits;
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+        }
         return hits;
     }
 

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

@@ -38,8 +38,10 @@
 import org.opensearch.search.SearchShardTarget;
 import org.opensearch.search.internal.ShardSearchContextId;
 import org.opensearch.search.query.QuerySearchResult;
+import org.opensearch.server.proto.QueryFetchSearchResultProto;
 
 import java.io.IOException;
+import java.io.InputStream;
 
 /**
  * Query fetch result
@@ -51,12 +53,20 @@
     private final QuerySearchResult queryResult;
     private final FetchSearchResult fetchResult;
 
+    private QueryFetchSearchResultProto.QueryFetchSearchResult queryFetchSearchResultProto;
+
     public QueryFetchSearchResult(StreamInput in) throws IOException {
         super(in);
         queryResult = new QuerySearchResult(in);
         fetchResult = new FetchSearchResult(in);
     }
 
+    public QueryFetchSearchResult(InputStream in) throws IOException {
+        this.queryFetchSearchResultProto = QueryFetchSearchResultProto.QueryFetchSearchResult.parseFrom(in);
+        queryResult = new QuerySearchResult(in);
+        fetchResult = new FetchSearchResult(in);
+    }
+
     public QueryFetchSearchResult(QuerySearchResult queryResult, FetchSearchResult fetchResult) {
         this.queryResult = queryResult;
         this.fetchResult = fetchResult;