Support derived fields definition in search request

* adds support for fetch phase on derived fields
* adds support for highligting on derived fields

Signed-off-by: Rishabh Maurya <rishabhmaurya05@gmail.com>

Author: rishabhmaurya

server/src/main/java/org/opensearch/index/mapper/DocumentMapperParser.java

@@ -133,7 +133,7 @@ public DocumentMapper parse(@Nullable String type, CompressedXContent source) th
     }
 
     @SuppressWarnings({ "unchecked" })
-    private DocumentMapper parse(String type, Map<String, Object> mapping) throws MapperParsingException {
+    public DocumentMapper parse(String type, Map<String, Object> mapping) throws MapperParsingException {
         if (type == null) {
             throw new MapperParsingException("Failed to derive type");
         }

server/src/main/java/org/opensearch/index/query/QueryShardContext.java

@@ -58,6 +58,7 @@
 import org.opensearch.index.cache.bitset.BitsetFilterCache;
 import org.opensearch.index.fielddata.IndexFieldData;
 import org.opensearch.index.mapper.ContentPath;
+import org.opensearch.index.mapper.DerivedFieldMapper;
 import org.opensearch.index.mapper.DocumentMapper;
 import org.opensearch.index.mapper.MappedFieldType;
 import org.opensearch.index.mapper.Mapper;
@@ -119,6 +120,8 @@ public class QueryShardContext extends QueryRewriteContext {
     private final ValuesSourceRegistry valuesSourceRegistry;
     private BitSetProducer parentFilter;
 
+    private DocumentMapper derivedFieldMappers;
+
     public QueryShardContext(
         int shardId,
         IndexSettings indexSettings,
@@ -264,6 +267,7 @@ private QueryShardContext(
         this.fullyQualifiedIndex = fullyQualifiedIndex;
         this.allowExpensiveQueries = allowExpensiveQueries;
         this.valuesSourceRegistry = valuesSourceRegistry;
+        derivedFieldMappers = null;
     }
 
     private void reset() {
@@ -395,6 +399,14 @@ public ValuesSourceRegistry getValuesSourceRegistry() {
         return valuesSourceRegistry;
     }
 
+    public void setDerivedFieldMappers(DocumentMapper derivedFieldMappers) {
+        this.derivedFieldMappers = derivedFieldMappers;
+    }
+
+    public DocumentMapper getDerivedFieldsMapper() {
+        return derivedFieldMappers;
+    }
+
     public void setAllowUnmappedFields(boolean allowUnmappedFields) {
         this.allowUnmappedFields = allowUnmappedFields;
     }
@@ -404,14 +416,20 @@ public void setMapUnmappedFieldAsString(boolean mapUnmappedFieldAsString) {
     }
 
     MappedFieldType failIfFieldMappingNotFound(String name, MappedFieldType fieldMapping) {
-        if (fieldMapping != null || allowUnmappedFields) {
+        if (fieldMapping != null) {
             return fieldMapping;
-        } else if (mapUnmappedFieldAsString) {
-            TextFieldMapper.Builder builder = new TextFieldMapper.Builder(name, mapperService.getIndexAnalyzers());
-            return builder.build(new Mapper.BuilderContext(indexSettings.getSettings(), new ContentPath(1))).fieldType();
-        } else {
-            throw new QueryShardException(this, "No field mapping can be found for the field with name [{}]", name);
-        }
+        } else if (derivedFieldMappers != null
+            && derivedFieldMappers.mappers() != null
+            && derivedFieldMappers.mappers().getMapper(name) != null) {
+                return ((DerivedFieldMapper) derivedFieldMappers.mappers().getMapper(name)).fieldType();
+            } else if (allowUnmappedFields) {
+                return fieldMapping;
+            } else if (mapUnmappedFieldAsString) {
+                TextFieldMapper.Builder builder = new TextFieldMapper.Builder(name, mapperService.getIndexAnalyzers());
+                return builder.build(new Mapper.BuilderContext(indexSettings.getSettings(), new ContentPath(1))).fieldType();
+            } else {
+                throw new QueryShardException(this, "No field mapping can be found for the field with name [{}]", name);
+            }
     }
 
     private SearchLookup lookup = null;

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

@@ -77,6 +77,8 @@
 import org.opensearch.index.IndexService;
 import org.opensearch.index.IndexSettings;
 import org.opensearch.index.engine.Engine;
+import org.opensearch.index.mapper.DerivedFieldMapper;
+import org.opensearch.index.mapper.DocumentMapper;
 import org.opensearch.index.query.InnerHitContextBuilder;
 import org.opensearch.index.query.MatchAllQueryBuilder;
 import org.opensearch.index.query.MatchNoneQueryBuilder;
@@ -1066,6 +1068,15 @@ private DefaultSearchContext createSearchContext(ReaderContext reader, ShardSear
             // might end up with incorrect state since we are using now() or script services
             // during rewrite and normalized / evaluate templates etc.
             QueryShardContext context = new QueryShardContext(searchContext.getQueryShardContext());
+            if (request.source().derivedFields() != null && request.source().size() != 0) {
+                Map<String, Object> derivedFieldObject = new HashMap<>();
+                derivedFieldObject.put(DerivedFieldMapper.CONTENT_TYPE, request.source().derivedFields());
+                DocumentMapper documentMapper = searchContext.mapperService()
+                    .documentMapperParser()
+                    .parse(DerivedFieldMapper.CONTENT_TYPE, derivedFieldObject);
+                context.setDerivedFieldMappers(documentMapper);
+                searchContext.getQueryShardContext().setDerivedFieldMappers(documentMapper);
+            }
             Rewriteable.rewrite(request.getRewriteable(), context, true);
             assert searchContext.getQueryShardContext().isCacheable();
             success = true;

server/src/main/java/org/opensearch/search/builder/SearchSourceBuilder.java

@@ -51,6 +51,7 @@
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.core.xcontent.XContentHelper;
 import org.opensearch.core.xcontent.XContentParser;
+import org.opensearch.index.mapper.DerivedFieldMapper;
 import org.opensearch.index.query.QueryBuilder;
 import org.opensearch.index.query.QueryRewriteContext;
 import org.opensearch.index.query.Rewriteable;
@@ -113,6 +114,7 @@ public final class SearchSourceBuilder implements Writeable, ToXContentObject, R
     public static final ParseField DOCVALUE_FIELDS_FIELD = new ParseField("docvalue_fields");
     public static final ParseField FETCH_FIELDS_FIELD = new ParseField("fields");
     public static final ParseField SCRIPT_FIELDS_FIELD = new ParseField("script_fields");
+    public static final ParseField DERIVED_FIELDS_FIELD = new ParseField(DerivedFieldMapper.CONTENT_TYPE);
     public static final ParseField SCRIPT_FIELD = new ParseField("script");
     public static final ParseField IGNORE_FAILURE_FIELD = new ParseField("ignore_failure");
     public static final ParseField SORT_FIELD = new ParseField("sort");
@@ -192,6 +194,7 @@ public static HighlightBuilder highlight() {
     private StoredFieldsContext storedFieldsContext;
     private List<FieldAndFormat> docValueFields;
     private List<ScriptField> scriptFields;
+    private Map<String, Object> derivedFields;
     private FetchSourceContext fetchSourceContext;
     private List<FieldAndFormat> fetchFields;
 
@@ -247,6 +250,9 @@ public SearchSourceBuilder(StreamInput in) throws IOException {
         if (in.readBoolean()) {
             scriptFields = in.readList(ScriptField::new);
         }
+        if (in.readBoolean()) {
+            derivedFields = in.readMap();
+        }
         size = in.readVInt();
         if (in.readBoolean()) {
             int size = in.readVInt();
@@ -310,6 +316,11 @@ public void writeTo(StreamOutput out) throws IOException {
         if (hasScriptFields) {
             out.writeList(scriptFields);
         }
+        boolean hasDerivedFields = derivedFields != null;
+        out.writeBoolean(hasDerivedFields);
+        if (hasDerivedFields) {
+            out.writeMap(derivedFields);
+        }
         out.writeVInt(size);
         boolean hasSorts = sorts != null;
         out.writeBoolean(hasSorts);
@@ -955,6 +966,10 @@ public List<ScriptField> scriptFields() {
         return scriptFields;
     }
 
+    public Map<String, Object> derivedFields() {
+        return derivedFields;
+    }
+
     /**
      * Sets the boost a specific index or alias will receive when the query is executed
      * against it.
@@ -1119,6 +1134,7 @@ private SearchSourceBuilder shallowCopy(
         rewrittenBuilder.queryBuilder = queryBuilder;
         rewrittenBuilder.rescoreBuilders = rescoreBuilders;
         rewrittenBuilder.scriptFields = scriptFields;
+        rewrittenBuilder.derivedFields = derivedFields;
         rewrittenBuilder.searchAfterBuilder = searchAfterBuilder;
         rewrittenBuilder.sliceBuilder = slice;
         rewrittenBuilder.size = size;
@@ -1220,6 +1236,8 @@ public void parseXContent(XContentParser parser, boolean checkTrailingTokens) th
                     while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
                         scriptFields.add(new ScriptField(parser));
                     }
+                } else if (DERIVED_FIELDS_FIELD.match(currentFieldName, parser.getDeprecationHandler())) {
+                    derivedFields = parser.map();
                 } else if (INDICES_BOOST_FIELD.match(currentFieldName, parser.getDeprecationHandler())) {
                     deprecationLogger.deprecate(
                         "indices_boost_object_format",
@@ -1434,6 +1452,10 @@ public XContentBuilder innerToXContent(XContentBuilder builder, Params params) t
             builder.endObject();
         }
 
+        if (derivedFields != null) {
+            builder.field(DERIVED_FIELDS_FIELD.getPreferredName(), derivedFields);
+        }
+
         if (sorts != null) {
             builder.startArray(SORT_FIELD.getPreferredName());
             for (SortBuilder<?> sort : sorts) {
@@ -1772,6 +1794,7 @@ public int hashCode() {
             queryBuilder,
             rescoreBuilders,
             scriptFields,
+            derivedFields,
             size,
             sorts,
             searchAfterBuilder,
@@ -1815,6 +1838,7 @@ public boolean equals(Object obj) {
             && Objects.equals(queryBuilder, other.queryBuilder)
             && Objects.equals(rescoreBuilders, other.rescoreBuilders)
             && Objects.equals(scriptFields, other.scriptFields)
+            && Objects.equals(derivedFields, other.derivedFields)
             && Objects.equals(size, other.size)
             && Objects.equals(sorts, other.sorts)
             && Objects.equals(searchAfterBuilder, other.searchAfterBuilder)

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

@@ -35,6 +35,8 @@
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.search.Query;
 import org.opensearch.common.regex.Regex;
+import org.opensearch.index.mapper.DerivedFieldMapper;
+import org.opensearch.index.mapper.DocumentMapper;
 import org.opensearch.index.mapper.KeywordFieldMapper;
 import org.opensearch.index.mapper.MappedFieldType;
 import org.opensearch.index.mapper.SourceFieldMapper;
@@ -145,6 +147,12 @@ private Map<String, Function<HitContext, FieldHighlightContext>> contextBuilders
             boolean fieldNameContainsWildcards = field.field().contains("*");
             for (String fieldName : fieldNamesToHighlight) {
                 MappedFieldType fieldType = context.mapperService().fieldType(fieldName);
+                if (fieldType == null && context.getQueryShardContext().getDerivedFieldsMapper() != null) {
+                    DocumentMapper derivedFieldsMapper = context.getQueryShardContext().getDerivedFieldsMapper();
+                    if (derivedFieldsMapper.mappers() != null && derivedFieldsMapper.mappers().getMapper(fieldName) != null) {
+                        fieldType = ((DerivedFieldMapper) derivedFieldsMapper.mappers().getMapper(fieldName)).fieldType();
+                    }
+                }
                 if (fieldType == null) {
                     continue;
                 }
@@ -170,12 +178,13 @@ private Map<String, Function<HitContext, FieldHighlightContext>> contextBuilders
                 Query highlightQuery = field.fieldOptions().highlightQuery();
 
                 boolean forceSource = highlightContext.forceSource(field);
+                MappedFieldType finalFieldType = fieldType;
                 builders.put(
                     fieldName,
                     hc -> new FieldHighlightContext(
-                        fieldType.name(),
+                        finalFieldType.name(),
                         field,
-                        fieldType,
+                        finalFieldType,
                         context,
                         hc,
                         highlightQuery == null ? query : highlightQuery,

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

@@ -35,6 +35,7 @@
 import org.apache.lucene.search.highlight.Encoder;
 import org.apache.lucene.search.highlight.SimpleHTMLEncoder;
 import org.opensearch.index.fieldvisitor.CustomFieldsVisitor;
+import org.opensearch.index.mapper.DerivedFieldValueFetcher;
 import org.opensearch.index.mapper.MappedFieldType;
 import org.opensearch.index.mapper.ValueFetcher;
 import org.opensearch.index.query.QueryShardContext;
@@ -77,6 +78,9 @@ public static List<Object> loadFieldValues(
             return textsToHighlight != null ? textsToHighlight : Collections.emptyList();
         }
         ValueFetcher fetcher = fieldType.valueFetcher(context, null, null);
+        if (fetcher instanceof DerivedFieldValueFetcher) {
+            fetcher.setNextReader(hitContext.reader().getContext());
+        }
         return fetcher.fetchValues(hitContext.sourceLookup());
     }
 

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

@@ -159,6 +159,12 @@ CustomUnifiedHighlighter buildHighlighter(FieldHighlightContext fieldContext) th
         Integer fieldMaxAnalyzedOffset = fieldContext.field.fieldOptions().maxAnalyzerOffset();
         int numberOfFragments = fieldContext.field.fieldOptions().numberOfFragments();
         Analyzer analyzer = getAnalyzer(fieldContext.context.mapperService().documentMapper());
+        if (fieldContext.context.getQueryShardContext().getDerivedFieldsMapper() != null) {
+            DocumentMapper derivedFieldsMapper = fieldContext.context.getQueryShardContext().getDerivedFieldsMapper();
+            if (derivedFieldsMapper != null) {
+                analyzer = getAnalyzer(derivedFieldsMapper);
+            }
+        }
         if (fieldMaxAnalyzedOffset != null) {
             analyzer = getLimitedOffsetAnalyzer(analyzer, fieldMaxAnalyzedOffset);
         }