diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/AutocompleteRequestHandler.java b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/AutocompleteRequestHandler.java index 24598d7fdaadb..b7a04f2064d9b 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/AutocompleteRequestHandler.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/AutocompleteRequestHandler.java @@ -140,10 +140,15 @@ public SearchRequest getSearchRequest( opContext.getObjectMapper(), cac, queryWithDefaultFilters, - customQueryConfig)) + customQueryConfig, + input)) .orElse( SearchQueryBuilder.buildScoreFunctions( - opContext, customQueryConfig, List.of(entitySpec), queryWithDefaultFilters)); + opContext, + customQueryConfig, + List.of(entitySpec), + input, + queryWithDefaultFilters)); searchSourceBuilder.query(functionScoreQueryBuilder); ESUtils.buildSortOrder(searchSourceBuilder, null, List.of(entitySpec)); diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/CustomizedQueryHandler.java b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/CustomizedQueryHandler.java index 0dbdf80860f7f..c4418401a844e 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/CustomizedQueryHandler.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/CustomizedQueryHandler.java @@ -119,16 +119,18 @@ private static BoolQueryBuilder toBoolQueryBuilder( public static FunctionScoreQueryBuilder functionScoreQueryBuilder( @Nonnull ObjectMapper objectMapper, @Nonnull QueryConfiguration customQueryConfiguration, - QueryBuilder queryBuilder) { + QueryBuilder queryBuilder, + String query) { return toFunctionScoreQueryBuilder( - objectMapper, queryBuilder, customQueryConfiguration.getFunctionScore()); + objectMapper, queryBuilder, customQueryConfiguration.getFunctionScore(), query); } public static Optional functionScoreQueryBuilder( @Nonnull ObjectMapper objectMapper, @Nonnull AutocompleteConfiguration customAutocompleteConfiguration, QueryBuilder queryBuilder, - @Nullable QueryConfiguration customQueryConfiguration) { + @Nullable QueryConfiguration customQueryConfiguration, + String query) { Optional result = Optional.empty(); @@ -143,14 +145,17 @@ public static Optional functionScoreQueryBuilder( result = Optional.of( toFunctionScoreQueryBuilder( - objectMapper, queryBuilder, customQueryConfiguration.getFunctionScore())); + objectMapper, queryBuilder, customQueryConfiguration.getFunctionScore(), query)); } else if (customAutocompleteConfiguration.getFunctionScore() != null && !customAutocompleteConfiguration.getFunctionScore().isEmpty()) { log.debug("Applying custom autocomplete function scores."); result = Optional.of( toFunctionScoreQueryBuilder( - objectMapper, queryBuilder, customAutocompleteConfiguration.getFunctionScore())); + objectMapper, + queryBuilder, + customAutocompleteConfiguration.getFunctionScore(), + query)); } return result; @@ -159,7 +164,8 @@ public static Optional functionScoreQueryBuilder( private static FunctionScoreQueryBuilder toFunctionScoreQueryBuilder( @Nonnull ObjectMapper objectMapper, @Nonnull QueryBuilder queryBuilder, - @Nonnull Map params) { + @Nonnull Map params, + String query) { try { HashMap body = new HashMap<>(params); if (!body.isEmpty()) { @@ -168,7 +174,12 @@ private static FunctionScoreQueryBuilder toFunctionScoreQueryBuilder( body.put("query", objectMapper.readValue(queryBuilder.toString(), Map.class)); - String jsonFragment = objectMapper.writeValueAsString(Map.of("function_score", body)); + String jsonFragment = + objectMapper + .writeValueAsString(Map.of("function_score", body)) + .replace("\"{{query_string}}\"", objectMapper.writeValueAsString(query)) + .replace( + "\"{{unquoted_query_string}}\"", objectMapper.writeValueAsString(unquote(query))); XContentParser parser = XContentType.JSON .xContent() diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchQueryBuilder.java b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchQueryBuilder.java index 9ac460ed7e1e1..3e76d3600d6a6 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchQueryBuilder.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchQueryBuilder.java @@ -78,7 +78,7 @@ public QueryBuilder buildQuery( final QueryBuilder queryBuilder = buildInternalQuery(opContext, customQueryConfig, entitySpecs, query, fulltext); - return buildScoreFunctions(opContext, customQueryConfig, entitySpecs, queryBuilder); + return buildScoreFunctions(opContext, customQueryConfig, entitySpecs, query, queryBuilder); } /** @@ -485,12 +485,13 @@ static FunctionScoreQueryBuilder buildScoreFunctions( @Nonnull OperationContext opContext, @Nullable QueryConfiguration customQueryConfig, @Nonnull List entitySpecs, + String query, @Nonnull QueryBuilder queryBuilder) { if (customQueryConfig != null) { // Prefer configuration function scoring over annotation scoring return CustomizedQueryHandler.functionScoreQueryBuilder( - opContext.getObjectMapper(), customQueryConfig, queryBuilder); + opContext.getObjectMapper(), customQueryConfig, queryBuilder, query); } else { return QueryBuilders.functionScoreQuery( queryBuilder, buildAnnotationScoreFunctions(entitySpecs)) diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/CustomizedQueryHandlerTest.java b/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/CustomizedQueryHandlerTest.java index 4e4c8acf300e4..a16e497ff548f 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/CustomizedQueryHandlerTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/CustomizedQueryHandlerTest.java @@ -179,7 +179,7 @@ public void functionScoreQueryBuilderTest() { */ FunctionScoreQueryBuilder selectStarTest = CustomizedQueryHandler.functionScoreQueryBuilder( - new ObjectMapper(), test.lookupQueryConfig("*").get(), inputQuery); + new ObjectMapper(), test.lookupQueryConfig("*").get(), inputQuery, "*"); FunctionScoreQueryBuilder.FilterFunctionBuilder[] expectedSelectStarScoreFunctions = { new FunctionScoreQueryBuilder.FilterFunctionBuilder( @@ -203,7 +203,7 @@ public void functionScoreQueryBuilderTest() { */ FunctionScoreQueryBuilder defaultTest = CustomizedQueryHandler.functionScoreQueryBuilder( - new ObjectMapper(), test.lookupQueryConfig("foobar").get(), inputQuery); + new ObjectMapper(), test.lookupQueryConfig("foobar").get(), inputQuery, "foobar"); FunctionScoreQueryBuilder.FilterFunctionBuilder[] expectedDefaultScoreFunctions = { new FunctionScoreQueryBuilder.FilterFunctionBuilder(