Feature/custom highlight on search (#11339)

Co-authored-by: Arpan Chakraborty <arpan.chakraborty@blackrock.com>
Co-authored-by: krekacsk <krekacs.k@outlook.com>
This commit is contained in:
arpanchakra29 2024-09-11 16:25:48 +01:00 committed by GitHub
parent 728a41e895
commit 33e898af4a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 77 additions and 1 deletions

View File

@ -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 class SearchFlagsInputMapper
.map(c -> GroupingCriterionInputMapper.map(context, c))
.collect(Collectors.toList()))));
}
if (searchFlags.getCustomHighlightingFields() != null) {
result.setCustomHighlightingFields(
new StringArray(searchFlags.getCustomHighlightingFields()));
}
return result;
}
}

View File

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

View File

@ -49,6 +49,7 @@ import java.util.stream.Stream;
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 class SearchRequestHandler {
.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 @@ public class SearchRequestHandler {
}
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;
}
}

View File

@ -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 class SearchRequestHandlerTest extends AbstractTestNGSpringContextTests {
"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 =

View File

@ -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]
}

View File

@ -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",