Feature/custom highlight on search

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/common/mappers/SearchFlagsInputMapper.java

@@ -1,5 +1,6 @@
 package com.linkedin.datahub.graphql.types.common.mappers;
 
+import com.linkedin.data.template.StringArray;
 import com.linkedin.datahub.graphql.QueryContext;
 import com.linkedin.datahub.graphql.generated.SearchFlags;
 import com.linkedin.datahub.graphql.types.mappers.ModelMapper;
@@ -64,6 +65,10 @@ public com.linkedin.metadata.query.SearchFlags apply(
                           .map(c -> GroupingCriterionInputMapper.map(context, c))
                           .collect(Collectors.toList()))));
     }
+    if (searchFlags.getCustomHighlightingFields() != null) {
+      result.setCustomHighlightingFields(
+          new StringArray(searchFlags.getCustomHighlightingFields()));
+    }
     return result;
   }
 }

datahub-graphql-core/src/main/resources/search.graphql

@@ -162,6 +162,11 @@ input  SearchFlags {
   Whether to include restricted entities
   """
   includeRestricted: Boolean
+
+  """
+  fields to include for custom Highlighting
+  """
+  customHighlightingFields: [String!]
 }
 
 """

metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchRequestHandler.java

@@ -49,6 +49,7 @@
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 import lombok.extern.slf4j.Slf4j;
+import org.apache.commons.collections.CollectionUtils;
 import org.opensearch.action.search.SearchRequest;
 import org.opensearch.action.search.SearchResponse;
 import org.opensearch.common.unit.TimeValue;
@@ -211,8 +212,14 @@ public SearchRequest getSearchRequest(
           .forEach(searchSourceBuilder::aggregation);
     }
     if (Boolean.FALSE.equals(searchFlags.isSkipHighlighting())) {
-      searchSourceBuilder.highlighter(highlights);
+      if (CollectionUtils.isNotEmpty(searchFlags.getCustomHighlightingFields())) {
+        searchSourceBuilder.highlighter(
+            getValidatedHighlighter(searchFlags.getCustomHighlightingFields()));
+      } else {
+        searchSourceBuilder.highlighter(highlights);
+      }
     }
+
     ESUtils.buildSortOrder(searchSourceBuilder, sortCriteria, entitySpecs);
 
     if (Boolean.TRUE.equals(searchFlags.isGetSuggestions())) {
@@ -556,4 +563,16 @@ private List<SearchSuggestion> extractSearchSuggestions(@Nonnull SearchResponse
     }
     return searchSuggestions;
   }
+
+  private HighlightBuilder getValidatedHighlighter(Collection<String> fieldsToHighlight) {
+    HighlightBuilder highlightBuilder = new HighlightBuilder();
+    highlightBuilder.preTags("");
+    highlightBuilder.postTags("");
+    fieldsToHighlight.stream()
+        .filter(defaultQueryFieldNames::contains)
+        .flatMap(fieldName -> Stream.of(fieldName, fieldName + ".*"))
+        .distinct()
+        .forEach(highlightBuilder::field);
+    return highlightBuilder;
+  }
 }

metadata-io/src/test/java/com/linkedin/metadata/search/query/request/SearchRequestHandlerTest.java

@@ -55,6 +55,14 @@ public class SearchRequestHandlerTest extends AbstractTestNGSpringContextTests {
   private OperationContext operationContext;
 
   public static SearchConfiguration testQueryConfig;
+  public static List<String> validHighlightingFields = List.of("urn", "foreignKey");
+  public static StringArray customHighlightFields =
+      new StringArray(
+          List.of(
+              validHighlightingFields.get(0),
+              validHighlightingFields.get(1),
+              "notExistingField",
+              ""));
 
   static {
     testQueryConfig = new SearchConfiguration();
@@ -102,6 +110,32 @@ public void testDatasetFieldsAndHighlights() {
         "unexpected lineage fields in highlights: " + highlightFields);
   }
 
+  @Test
+  public void testCustomHighlights() {
+    EntitySpec entitySpec = operationContext.getEntityRegistry().getEntitySpec("dataset");
+    SearchRequestHandler requestHandler =
+        SearchRequestHandler.getBuilder(TestEntitySpecBuilder.getSpec(), testQueryConfig, null);
+    SearchRequest searchRequest =
+        requestHandler.getSearchRequest(
+            operationContext.withSearchFlags(
+                flags ->
+                    flags.setFulltext(false).setCustomHighlightingFields(customHighlightFields)),
+            "testQuery",
+            null,
+            null,
+            0,
+            10,
+            null);
+    SearchSourceBuilder sourceBuilder = searchRequest.source();
+    assertNotNull(sourceBuilder.highlighter());
+    assertEquals(4, sourceBuilder.highlighter().fields().size());
+    assertTrue(
+        sourceBuilder.highlighter().fields().stream()
+            .map(HighlightBuilder.Field::name)
+            .toList()
+            .containsAll(validHighlightingFields));
+  }
+
   @Test
   public void testSearchRequestHandlerHighlightingTurnedOff() {
     SearchRequestHandler requestHandler =

metadata-models/src/main/pegasus/com/linkedin/metadata/query/SearchFlags.pdl

@@ -48,4 +48,9 @@ record SearchFlags {
    * include restricted entities in results (default is to filter)
    */
    includeRestricted:optional boolean = false
+
+  /**
+   * Include mentioned fields inside elastic highlighting query
+   */
+   customHighlightingFields:optional array[string]
 }

metadata-service/restli-api/src/main/snapshot/com.linkedin.entity.aspects.snapshot.json

@@ -3180,7 +3180,7 @@
             }, {
               "name" : "isPartitioningKey",
               "type" : "boolean",
-              "doc" : "For Datasets which are partitioned, this determines the partitioning key.",
+              "doc" : "For Datasets which are partitioned, this determines the partitioning key.\nNote that multiple columns can be part of a partitioning key, but currently we do not support\nrendering the ordered partitioning key.",
               "optional" : true
             }, {
               "name" : "jsonProps",

metadata-service/restli-api/src/main/snapshot/com.linkedin.entity.entities.snapshot.json

@@ -3572,7 +3572,7 @@
                         }, {
                           "name" : "isPartitioningKey",
                           "type" : "boolean",
-                          "doc" : "For Datasets which are partitioned, this determines the partitioning key.",
+                          "doc" : "For Datasets which are partitioned, this determines the partitioning key.\nNote that multiple columns can be part of a partitioning key, but currently we do not support\nrendering the ordered partitioning key.",
                           "optional" : true
                         }, {
                           "name" : "jsonProps",
@@ -6032,6 +6032,14 @@
       "doc" : "include restricted entities in results (default is to filter)",
       "default" : false,
       "optional" : true
+    }, {
+      "name" : "customHighlightingFields",
+      "type" : {
+        "type" : "array",
+        "items" : "string"
+      },
+      "doc" : "Include mentioned fields inside elastic highlighting query",
+      "optional" : true
     } ]
   }, {
     "type" : "enum",

metadata-service/restli-api/src/main/snapshot/com.linkedin.entity.runs.snapshot.json

@@ -2908,7 +2908,7 @@
             }, {
               "name" : "isPartitioningKey",
               "type" : "boolean",
-              "doc" : "For Datasets which are partitioned, this determines the partitioning key.",
+              "doc" : "For Datasets which are partitioned, this determines the partitioning key.\nNote that multiple columns can be part of a partitioning key, but currently we do not support\nrendering the ordered partitioning key.",
               "optional" : true
             }, {
               "name" : "jsonProps",

metadata-service/restli-api/src/main/snapshot/com.linkedin.operations.operations.snapshot.json

@@ -2902,7 +2902,7 @@
             }, {
               "name" : "isPartitioningKey",
               "type" : "boolean",
-              "doc" : "For Datasets which are partitioned, this determines the partitioning key.",
+              "doc" : "For Datasets which are partitioned, this determines the partitioning key.\nNote that multiple columns can be part of a partitioning key, but currently we do not support\nrendering the ordered partitioning key.",
               "optional" : true
             }, {
               "name" : "jsonProps",

metadata-service/restli-api/src/main/snapshot/com.linkedin.platform.platform.snapshot.json

@@ -3566,7 +3566,7 @@
                         }, {
                           "name" : "isPartitioningKey",
                           "type" : "boolean",
-                          "doc" : "For Datasets which are partitioned, this determines the partitioning key.",
+                          "doc" : "For Datasets which are partitioned, this determines the partitioning key.\nNote that multiple columns can be part of a partitioning key, but currently we do not support\nrendering the ordered partitioning key.",
                           "optional" : true
                         }, {
                           "name" : "jsonProps",