Skip to content

Commit 52f71dd

Browse files
fix(autocomplete): fix autocomplete duplicate field (#12558)
1 parent 65376ee commit 52f71dd

File tree

8 files changed

+173
-132
lines changed

8 files changed

+173
-132
lines changed

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -642,7 +642,7 @@ private QueryBuilder buildQueryStringV2(
642642
EntitySpec entitySpec = opContext.getEntityRegistry().getEntitySpec(entityName);
643643
QueryBuilder query =
644644
SearchRequestHandler.getBuilder(
645-
opContext.getEntityRegistry(),
645+
opContext,
646646
entitySpec,
647647
searchConfiguration,
648648
customSearchConfiguration,
@@ -683,7 +683,7 @@ private QueryBuilder buildQueryStringBrowseAcrossEntities(
683683

684684
QueryBuilder query =
685685
SearchRequestHandler.getBuilder(
686-
finalOpContext.getEntityRegistry(),
686+
finalOpContext,
687687
entitySpecs,
688688
searchConfiguration,
689689
customSearchConfiguration,

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

Lines changed: 11 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -150,7 +150,7 @@ private SearchResult executeAndExtract(
150150
return transformIndexIntoEntityName(
151151
opContext.getSearchContext().getIndexConvention(),
152152
SearchRequestHandler.getBuilder(
153-
opContext.getEntityRegistry(),
153+
opContext,
154154
entitySpec,
155155
searchConfiguration,
156156
customSearchConfiguration,
@@ -257,7 +257,7 @@ private ScrollResult executeAndExtract(
257257
return transformIndexIntoEntityName(
258258
opContext.getSearchContext().getIndexConvention(),
259259
SearchRequestHandler.getBuilder(
260-
opContext.getEntityRegistry(),
260+
opContext,
261261
entitySpecs,
262262
searchConfiguration,
263263
customSearchConfiguration,
@@ -311,7 +311,7 @@ public SearchResult search(
311311
"searchRequest",
312312
() ->
313313
SearchRequestHandler.getBuilder(
314-
opContext.getEntityRegistry(),
314+
opContext,
315315
entitySpecs,
316316
searchConfiguration,
317317
customSearchConfiguration,
@@ -357,7 +357,7 @@ public SearchResult filter(
357357
Filter transformedFilters = transformFilterForEntities(filters, indexConvention);
358358
final SearchRequest searchRequest =
359359
SearchRequestHandler.getBuilder(
360-
opContext.getEntityRegistry(),
360+
opContext,
361361
entitySpec,
362362
searchConfiguration,
363363
customSearchConfiguration,
@@ -395,7 +395,11 @@ public AutoCompleteResult autoComplete(
395395
IndexConvention indexConvention = opContext.getSearchContext().getIndexConvention();
396396
AutocompleteRequestHandler builder =
397397
AutocompleteRequestHandler.getBuilder(
398-
entitySpec, customSearchConfiguration, queryFilterRewriteChain, searchConfiguration);
398+
opContext,
399+
entitySpec,
400+
customSearchConfiguration,
401+
queryFilterRewriteChain,
402+
searchConfiguration);
399403
SearchRequest req =
400404
builder.getSearchRequest(
401405
opContext,
@@ -441,7 +445,7 @@ public Map<String, Long> aggregateByValue(
441445
IndexConvention indexConvention = opContext.getSearchContext().getIndexConvention();
442446
final SearchRequest searchRequest =
443447
SearchRequestHandler.getBuilder(
444-
opContext.getEntityRegistry(),
448+
opContext,
445449
entitySpecs,
446450
searchConfiguration,
447451
customSearchConfiguration,
@@ -578,7 +582,7 @@ private SearchRequest getScrollRequest(
578582
}
579583

580584
return SearchRequestHandler.getBuilder(
581-
opContext.getEntityRegistry(),
585+
opContext,
582586
entitySpecs,
583587
searchConfiguration,
584588
customSearchConfiguration,

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

Lines changed: 31 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import io.datahubproject.metadata.context.OperationContext;
2424
import java.net.URISyntaxException;
2525
import java.util.ArrayList;
26+
import java.util.Collection;
2627
import java.util.Collections;
2728
import java.util.HashSet;
2829
import java.util.List;
@@ -45,9 +46,9 @@
4546
import org.opensearch.search.fetch.subphase.highlight.HighlightBuilder;
4647

4748
@Slf4j
48-
public class AutocompleteRequestHandler {
49+
public class AutocompleteRequestHandler extends BaseRequestHandler {
4950

50-
private final List<Pair> _defaultAutocompleteFields;
51+
private final List<Pair<String, String>> _defaultAutocompleteFields;
5152
private final Map<String, Set<SearchableAnnotation.FieldType>> searchableFieldTypes;
5253

5354
private static final Map<EntitySpec, AutocompleteRequestHandler>
@@ -58,8 +59,10 @@ public class AutocompleteRequestHandler {
5859
private final EntitySpec entitySpec;
5960
private final QueryFilterRewriteChain queryFilterRewriteChain;
6061
private final SearchConfiguration searchConfiguration;
62+
@Nonnull private final HighlightBuilder highlights;
6163

6264
public AutocompleteRequestHandler(
65+
@Nonnull OperationContext systemOperationContext,
6366
@Nonnull EntitySpec entitySpec,
6467
@Nullable CustomSearchConfiguration customSearchConfiguration,
6568
@Nonnull QueryFilterRewriteChain queryFilterRewriteChain,
@@ -79,6 +82,7 @@ public AutocompleteRequestHandler(
7982
Double.toString(searchableAnnotation.getBoostScore()))),
8083
Stream.of(Pair.of("urn", "1.0")))
8184
.collect(Collectors.toList());
85+
this.highlights = getDefaultHighlights(systemOperationContext);
8286
searchableFieldTypes =
8387
fieldSpecs.stream()
8488
.collect(
@@ -98,6 +102,7 @@ public AutocompleteRequestHandler(
98102
}
99103

100104
public static AutocompleteRequestHandler getBuilder(
105+
@Nonnull OperationContext systemOperationContext,
101106
@Nonnull EntitySpec entitySpec,
102107
@Nullable CustomSearchConfiguration customSearchConfiguration,
103108
@Nonnull QueryFilterRewriteChain queryFilterRewriteChain,
@@ -106,6 +111,7 @@ public static AutocompleteRequestHandler getBuilder(
106111
entitySpec,
107112
k ->
108113
new AutocompleteRequestHandler(
114+
systemOperationContext,
109115
entitySpec,
110116
customSearchConfiguration,
111117
queryFilterRewriteChain,
@@ -165,7 +171,8 @@ public SearchRequest getSearchRequest(
165171
ESUtils.buildSortOrder(searchSourceBuilder, null, List.of(entitySpec));
166172

167173
// wire inner non-scored query
168-
searchSourceBuilder.highlighter(getHighlights(field));
174+
searchSourceBuilder.highlighter(
175+
field == null || field.isEmpty() ? highlights : getHighlights(opContext, List.of(field)));
169176
searchRequest.source(searchSourceBuilder);
170177
return searchRequest;
171178
}
@@ -181,7 +188,7 @@ private BoolQueryBuilder getQuery(
181188
public BoolQueryBuilder getQuery(
182189
@Nonnull ObjectMapper objectMapper,
183190
@Nullable AutocompleteConfiguration customAutocompleteConfig,
184-
List<Pair> autocompleteFields,
191+
List<Pair<String, String>> autocompleteFields,
185192
@Nonnull String query) {
186193

187194
BoolQueryBuilder finalQuery =
@@ -201,7 +208,7 @@ public BoolQueryBuilder getQuery(
201208

202209
private Optional<QueryBuilder> getAutocompleteQuery(
203210
@Nullable AutocompleteConfiguration customConfig,
204-
List<Pair> autocompleteFields,
211+
List<Pair<String, String>> autocompleteFields,
205212
@Nonnull String query) {
206213
Optional<QueryBuilder> result = Optional.empty();
207214

@@ -212,7 +219,8 @@ private Optional<QueryBuilder> getAutocompleteQuery(
212219
return result;
213220
}
214221

215-
private BoolQueryBuilder defaultQuery(List<Pair> autocompleteFields, @Nonnull String query) {
222+
private BoolQueryBuilder defaultQuery(
223+
List<Pair<String, String>> autocompleteFields, @Nonnull String query) {
216224
BoolQueryBuilder finalQuery = QueryBuilders.boolQuery().minimumShouldMatch(1);
217225

218226
// Search for exact matches with higher boost and ngram matches
@@ -248,38 +256,25 @@ private BoolQueryBuilder defaultQuery(List<Pair> autocompleteFields, @Nonnull St
248256
return finalQuery;
249257
}
250258

251-
// Get HighlightBuilder to highlight the matched field
252-
private HighlightBuilder getHighlights(@Nullable String field) {
253-
HighlightBuilder highlightBuilder =
254-
new HighlightBuilder()
255-
// Don't set tags to get the original field value
256-
.preTags("")
257-
.postTags("")
258-
.numOfFragments(1);
259-
// Check for each field name and any subfields
260-
getAutocompleteFields(field)
261-
.forEach(
262-
pair -> {
263-
final String fieldName = (String) pair.getLeft();
264-
highlightBuilder
265-
.field(fieldName)
266-
.field(fieldName + ".*")
267-
.field(fieldName + ".ngram")
268-
.field(fieldName + ".delimited");
269-
if (!fieldName.equalsIgnoreCase("urn")) {
270-
highlightBuilder.field(fieldName + ".keyword");
271-
}
272-
});
273-
274-
// set field match req false for ngram
275-
highlightBuilder.fields().stream()
276-
.filter(f -> f.name().contains("ngram"))
277-
.forEach(f -> f.requireFieldMatch(false).noMatchSize(200));
278-
279-
return highlightBuilder;
259+
@Override
260+
public Collection<String> getDefaultQueryFieldNames() {
261+
return _defaultAutocompleteFields.stream().map(Pair::getKey).collect(Collectors.toList());
280262
}
281263

282-
private List<Pair> getAutocompleteFields(@Nullable String field) {
264+
@Override
265+
protected Collection<String> getValidQueryFieldNames() {
266+
return searchableFieldTypes.keySet();
267+
}
268+
269+
@Override
270+
protected Stream<String> highlightFieldExpansion(
271+
@Nonnull OperationContext opContext, @Nonnull String fieldName) {
272+
return Stream.concat(
273+
Stream.of(fieldName, fieldName + ".*", fieldName + ".ngram", fieldName + ".delimited"),
274+
Stream.of(ESUtils.toKeywordField(fieldName, false, opContext.getAspectRetriever())));
275+
}
276+
277+
private List<Pair<String, String>> getAutocompleteFields(@Nullable String field) {
283278
if (field != null && !field.isEmpty() && !field.equalsIgnoreCase("urn")) {
284279
return ImmutableList.of(Pair.of(field, "10.0"));
285280
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
package com.linkedin.metadata.search.elasticsearch.query.request;
2+
3+
import com.google.common.annotations.VisibleForTesting;
4+
import io.datahubproject.metadata.context.OperationContext;
5+
import java.util.Collection;
6+
import java.util.Objects;
7+
import java.util.stream.Stream;
8+
import javax.annotation.Nonnull;
9+
import javax.annotation.Nullable;
10+
import org.opensearch.search.fetch.subphase.highlight.HighlightBuilder;
11+
12+
public abstract class BaseRequestHandler {
13+
14+
/**
15+
* Provide the fields which are queried by default
16+
*
17+
* @return collection of field names
18+
*/
19+
protected abstract Collection<String> getDefaultQueryFieldNames();
20+
21+
protected abstract Collection<String> getValidQueryFieldNames();
22+
23+
protected abstract Stream<String> highlightFieldExpansion(
24+
@Nonnull OperationContext opContext, @Nonnull String fieldName);
25+
26+
@VisibleForTesting
27+
public HighlightBuilder getDefaultHighlights(@Nonnull OperationContext opContext) {
28+
return getHighlights(opContext, null);
29+
}
30+
31+
@VisibleForTesting
32+
public HighlightBuilder getHighlights(
33+
@Nonnull OperationContext opContext, @Nullable Collection<String> fieldsToHighlight) {
34+
HighlightBuilder highlightBuilder =
35+
new HighlightBuilder()
36+
// Don't set tags to get the original field value
37+
.preTags("")
38+
.postTags("")
39+
.numOfFragments(1);
40+
41+
final Stream<String> fieldStream;
42+
if (fieldsToHighlight == null || fieldsToHighlight.isEmpty()) {
43+
fieldStream = getDefaultQueryFieldNames().stream();
44+
} else {
45+
// filter for valid names
46+
fieldStream =
47+
fieldsToHighlight.stream()
48+
.filter(Objects::nonNull)
49+
.filter(fieldName -> !fieldName.isEmpty())
50+
.filter(getValidQueryFieldNames()::contains);
51+
}
52+
53+
fieldStream
54+
.flatMap(fieldName -> highlightFieldExpansion(opContext, fieldName))
55+
.distinct()
56+
.map(HighlightBuilder.Field::new)
57+
.map(
58+
field -> {
59+
if (field.name().endsWith("ngram")) {
60+
return field.requireFieldMatch(false).noMatchSize(200);
61+
} else {
62+
return field;
63+
}
64+
})
65+
.forEach(highlightBuilder::field);
66+
67+
return highlightBuilder;
68+
}
69+
}

0 commit comments

Comments
 (0)