Supply boosting to multi_match

Author: Samiul-TheSoccerFan

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticKnnVectorQueryRewriteInterceptor.java

@@ -47,8 +47,7 @@ protected String getQuery(QueryBuilder queryBuilder) {
         return queryVectorBuilder != null ? queryVectorBuilder.getModelText() : null;
     }
 
-    @Override
-    protected QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation) {
+    private QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation) {
         assert (queryBuilder instanceof KnnVectorQueryBuilder);
         KnnVectorQueryBuilder knnVectorQueryBuilder = (KnnVectorQueryBuilder) queryBuilder;
         Map<String, List<String>> inferenceIdsIndices = indexInformation.getInferenceIdsIndices();
@@ -68,6 +67,17 @@ protected QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceI
         return finalQueryBuilder;
     }
 
+    @Override
+    protected QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation, Float fieldWeight) {
+        QueryBuilder inferenceQuery = buildInferenceQuery(queryBuilder, indexInformation);
+
+        if (fieldWeight != null && fieldWeight.equals(1.0f) == false) {
+            inferenceQuery.boost(fieldWeight);
+        }
+
+        return inferenceQuery;
+    }
+
     private QueryBuilder buildInferenceQueryWithMultipleInferenceIds(
         KnnVectorQueryBuilder queryBuilder,
         Map<String, List<String>> inferenceIdsIndices

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticMatchQueryRewriteInterceptor.java

@@ -36,14 +36,24 @@ protected String getQuery(QueryBuilder queryBuilder) {
         return (String) matchQueryBuilder.value();
     }
 
-    @Override
-    protected QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation) {
+    private QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation) {
         SemanticQueryBuilder semanticQueryBuilder = new SemanticQueryBuilder(indexInformation.fieldName(), getQuery(queryBuilder), false);
         semanticQueryBuilder.boost(queryBuilder.boost());
         semanticQueryBuilder.queryName(queryBuilder.queryName());
         return semanticQueryBuilder;
     }
 
+    @Override
+    protected QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation, Float fieldWeight) {
+        QueryBuilder inferenceQuery = buildInferenceQuery(queryBuilder, indexInformation);
+
+        if (fieldWeight != null && fieldWeight.equals(1.0f) == false) {
+            inferenceQuery.boost(fieldWeight);
+        }
+
+        return inferenceQuery;
+    }
+
     @Override
     protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery(
         QueryBuilder queryBuilder,

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticMultiMatchQueryRewriteInterceptor.java

@@ -28,18 +28,28 @@ protected String getQuery(QueryBuilder queryBuilder) {
         return (String) multiMatchQueryBuilder.value();
     }
 
-    @Override
-    protected QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation) {
+    private QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation) {
         SemanticQueryBuilder semanticQueryBuilder = new SemanticQueryBuilder(
             indexInformation.fieldName(),
             getQuery(queryBuilder),
             false
         );
-        // TODO:: add boost
+        semanticQueryBuilder.boost(queryBuilder.boost());
         semanticQueryBuilder.queryName(queryBuilder.queryName());
         return semanticQueryBuilder;
     }
 
+    @Override
+    protected QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation, Float fieldWeight) {
+        QueryBuilder inferenceQuery = buildInferenceQuery(queryBuilder, indexInformation);
+
+        if (fieldWeight != null && fieldWeight.equals(1.0f) == false) {
+            inferenceQuery.boost(fieldWeight);
+        }
+
+        return inferenceQuery;
+    }
+
     @Override
     protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation) {
         assert (queryBuilder instanceof MultiMatchQueryBuilder);

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryRewriteInterceptor.java

@@ -42,7 +42,9 @@ public QueryBuilder interceptAndRewrite(QueryRewriteContext context, QueryBuilde
         }
 
         BoolQueryBuilder finalQueryBuilder = new BoolQueryBuilder();
-        for (String fieldName : fieldNamesWithWeights.keySet()) {
+        for (Map.Entry<String, Float> fieldSet : fieldNamesWithWeights.entrySet()) {
+            String fieldName = fieldSet.getKey();
+            Float fieldWeight = fieldSet.getValue();
             InferenceIndexInformationForField indexInformation = resolveIndicesForField(fieldName, resolvedIndices);
             if (indexInformation.getInferenceIndices().isEmpty()) {
                 // No inference fields were identified, so return the original query.
@@ -51,11 +53,11 @@ public QueryBuilder interceptAndRewrite(QueryRewriteContext context, QueryBuilde
                 // Combined case where the field name requested by this query contains both
                 // semantic_text and non-inference fields, so we have to combine queries per index
                 // containing each field type.
-                return finalQueryBuilder.should(buildCombinedInferenceAndNonInferenceQuery(queryBuilder, indexInformation));
+                finalQueryBuilder.should(buildCombinedInferenceAndNonInferenceQuery(queryBuilder, indexInformation));
             } else {
                 // The only fields we've identified are inference fields (e.g. semantic_text),
                 // so rewrite the entire query to work on a semantic_text field.
-                return finalQueryBuilder.should(buildInferenceQuery(queryBuilder, indexInformation));
+                finalQueryBuilder.should(buildInferenceQuery(queryBuilder, indexInformation, fieldWeight));
             }
         }
         return finalQueryBuilder;
@@ -79,9 +81,10 @@ public QueryBuilder interceptAndRewrite(QueryRewriteContext context, QueryBuilde
      *
      * @param queryBuilder {@link QueryBuilder}
      * @param indexInformation {@link InferenceIndexInformationForField}
+     * @param weight {@link Float}
      * @return {@link QueryBuilder}
      */
-    protected abstract QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation);
+    protected abstract QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation, Float weight);
 
     /**
      * Builds a combined inference and non-inference query,

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticSparseVectorQueryRewriteInterceptor.java

@@ -40,8 +40,7 @@ protected String getQuery(QueryBuilder queryBuilder) {
         return sparseVectorQueryBuilder.getQuery();
     }
 
-    @Override
-    protected QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation) {
+    private QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation) {
         Map<String, List<String>> inferenceIdsIndices = indexInformation.getInferenceIdsIndices();
         QueryBuilder finalQueryBuilder;
         if (inferenceIdsIndices.size() == 1) {
@@ -57,6 +56,17 @@ protected QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceI
         return finalQueryBuilder;
     }
 
+    @Override
+    protected QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation, Float fieldWeight) {
+        QueryBuilder inferenceQuery = buildInferenceQuery(queryBuilder, indexInformation);
+
+        if (fieldWeight != null && fieldWeight.equals(1.0f) == false) {
+            inferenceQuery.boost(fieldWeight);
+        }
+
+        return inferenceQuery;
+    }
+
     private QueryBuilder buildInferenceQueryWithMultipleInferenceIds(
         QueryBuilder queryBuilder,
         Map<String, List<String>> inferenceIdsIndices