fix(customSearch): apply query string interpolation to function score (…

…#11406)
datahub-project · Sep 17, 2024 · 5f53a49 · 5f53a49
1 parent c772a8c
commit 5f53a49
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 13 deletions.
diff --git a/.../com/linkedin/metadata/search/elasticsearch/query/request/AutocompleteRequestHandler.java b/.../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/...java/com/linkedin/metadata/search/elasticsearch/query/request/CustomizedQueryHandler.java b/...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> functionScoreQueryBuilder(
  @Nonnull ObjectMapper objectMapper,
  @Nonnull AutocompleteConfiguration customAutocompleteConfiguration,
  QueryBuilder queryBuilder,
- @Nullable QueryConfiguration customQueryConfiguration) {
+ @Nullable QueryConfiguration customQueryConfiguration,
+ String query) {
 
  Optional<FunctionScoreQueryBuilder> result = Optional.empty();
 
@@ -143,14 +145,17 @@ public static Optional<FunctionScoreQueryBuilder> 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> functionScoreQueryBuilder(
  private static FunctionScoreQueryBuilder toFunctionScoreQueryBuilder(
  @Nonnull ObjectMapper objectMapper,
  @Nonnull QueryBuilder queryBuilder,
- @Nonnull Map<String, Object> params) {
+ @Nonnull Map<String, Object> params,
+ String query) {
  try {
  HashMap<String, Object> 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/...ain/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchQueryBuilder.java b/...ain/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<EntitySpec> 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/.../src/test/java/com/linkedin/metadata/search/query/request/CustomizedQueryHandlerTest.java b/.../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(