23
23
import io .datahubproject .metadata .context .OperationContext ;
24
24
import java .net .URISyntaxException ;
25
25
import java .util .ArrayList ;
26
+ import java .util .Collection ;
26
27
import java .util .Collections ;
27
28
import java .util .HashSet ;
28
29
import java .util .List ;
45
46
import org .opensearch .search .fetch .subphase .highlight .HighlightBuilder ;
46
47
47
48
@ Slf4j
48
- public class AutocompleteRequestHandler {
49
+ public class AutocompleteRequestHandler extends BaseRequestHandler {
49
50
50
- private final List <Pair > _defaultAutocompleteFields ;
51
+ private final List <Pair < String , String > > _defaultAutocompleteFields ;
51
52
private final Map <String , Set <SearchableAnnotation .FieldType >> searchableFieldTypes ;
52
53
53
54
private static final Map <EntitySpec , AutocompleteRequestHandler >
@@ -58,8 +59,10 @@ public class AutocompleteRequestHandler {
58
59
private final EntitySpec entitySpec ;
59
60
private final QueryFilterRewriteChain queryFilterRewriteChain ;
60
61
private final SearchConfiguration searchConfiguration ;
62
+ @ Nonnull private final HighlightBuilder highlights ;
61
63
62
64
public AutocompleteRequestHandler (
65
+ @ Nonnull OperationContext systemOperationContext ,
63
66
@ Nonnull EntitySpec entitySpec ,
64
67
@ Nullable CustomSearchConfiguration customSearchConfiguration ,
65
68
@ Nonnull QueryFilterRewriteChain queryFilterRewriteChain ,
@@ -79,6 +82,7 @@ public AutocompleteRequestHandler(
79
82
Double .toString (searchableAnnotation .getBoostScore ()))),
80
83
Stream .of (Pair .of ("urn" , "1.0" )))
81
84
.collect (Collectors .toList ());
85
+ this .highlights = getDefaultHighlights (systemOperationContext );
82
86
searchableFieldTypes =
83
87
fieldSpecs .stream ()
84
88
.collect (
@@ -98,6 +102,7 @@ public AutocompleteRequestHandler(
98
102
}
99
103
100
104
public static AutocompleteRequestHandler getBuilder (
105
+ @ Nonnull OperationContext systemOperationContext ,
101
106
@ Nonnull EntitySpec entitySpec ,
102
107
@ Nullable CustomSearchConfiguration customSearchConfiguration ,
103
108
@ Nonnull QueryFilterRewriteChain queryFilterRewriteChain ,
@@ -106,6 +111,7 @@ public static AutocompleteRequestHandler getBuilder(
106
111
entitySpec ,
107
112
k ->
108
113
new AutocompleteRequestHandler (
114
+ systemOperationContext ,
109
115
entitySpec ,
110
116
customSearchConfiguration ,
111
117
queryFilterRewriteChain ,
@@ -165,7 +171,8 @@ public SearchRequest getSearchRequest(
165
171
ESUtils .buildSortOrder (searchSourceBuilder , null , List .of (entitySpec ));
166
172
167
173
// wire inner non-scored query
168
- searchSourceBuilder .highlighter (getHighlights (field ));
174
+ searchSourceBuilder .highlighter (
175
+ field == null || field .isEmpty () ? highlights : getHighlights (opContext , List .of (field )));
169
176
searchRequest .source (searchSourceBuilder );
170
177
return searchRequest ;
171
178
}
@@ -181,7 +188,7 @@ private BoolQueryBuilder getQuery(
181
188
public BoolQueryBuilder getQuery (
182
189
@ Nonnull ObjectMapper objectMapper ,
183
190
@ Nullable AutocompleteConfiguration customAutocompleteConfig ,
184
- List <Pair > autocompleteFields ,
191
+ List <Pair < String , String > > autocompleteFields ,
185
192
@ Nonnull String query ) {
186
193
187
194
BoolQueryBuilder finalQuery =
@@ -201,7 +208,7 @@ public BoolQueryBuilder getQuery(
201
208
202
209
private Optional <QueryBuilder > getAutocompleteQuery (
203
210
@ Nullable AutocompleteConfiguration customConfig ,
204
- List <Pair > autocompleteFields ,
211
+ List <Pair < String , String > > autocompleteFields ,
205
212
@ Nonnull String query ) {
206
213
Optional <QueryBuilder > result = Optional .empty ();
207
214
@@ -212,7 +219,8 @@ private Optional<QueryBuilder> getAutocompleteQuery(
212
219
return result ;
213
220
}
214
221
215
- private BoolQueryBuilder defaultQuery (List <Pair > autocompleteFields , @ Nonnull String query ) {
222
+ private BoolQueryBuilder defaultQuery (
223
+ List <Pair <String , String >> autocompleteFields , @ Nonnull String query ) {
216
224
BoolQueryBuilder finalQuery = QueryBuilders .boolQuery ().minimumShouldMatch (1 );
217
225
218
226
// Search for exact matches with higher boost and ngram matches
@@ -248,38 +256,25 @@ private BoolQueryBuilder defaultQuery(List<Pair> autocompleteFields, @Nonnull St
248
256
return finalQuery ;
249
257
}
250
258
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 ());
280
262
}
281
263
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 ) {
283
278
if (field != null && !field .isEmpty () && !field .equalsIgnoreCase ("urn" )) {
284
279
return ImmutableList .of (Pair .of (field , "10.0" ));
285
280
}
0 commit comments