working commit

Author: Samiul-TheSoccerFan

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

@@ -94,8 +94,7 @@ private QueryBuilder buildInferenceQueryWithMultipleInferenceIds(
         return boolQueryBuilder;
     }
 
-    @Override
-    protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery(
+    private QueryBuilder buildCombinedInferenceAndNonInferenceQuery(
         QueryBuilder queryBuilder,
         InferenceIndexInformationForField indexInformation
     ) {
@@ -121,6 +120,21 @@ protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery(
         return boolQueryBuilder;
     }
 
+    @Override
+    protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery(
+        QueryBuilder queryBuilder,
+        InferenceIndexInformationForField indexInformation,
+        Float fieldWeight
+    ) {
+        QueryBuilder inferenceQuery = buildCombinedInferenceAndNonInferenceQuery(queryBuilder, indexInformation);
+
+        if (fieldWeight != null && fieldWeight.equals(1.0f) == false) {
+            inferenceQuery.boost(fieldWeight);
+        }
+
+        return inferenceQuery;
+    }
+
     private QueryBuilder buildNestedQueryFromKnnVectorQuery(
         KnnVectorQueryBuilder knnVectorQueryBuilder,
         List<String> indices,

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

@@ -54,7 +54,6 @@ protected QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceI
         return inferenceQuery;
     }
 
-    @Override
     protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery(
         QueryBuilder queryBuilder,
         InferenceIndexInformationForField indexInformation
@@ -78,6 +77,21 @@ protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery(
         return boolQueryBuilder;
     }
 
+    @Override
+    protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery(
+        QueryBuilder queryBuilder,
+        InferenceIndexInformationForField indexInformation,
+        Float fieldWeight
+    ) {
+        QueryBuilder inferenceQuery = buildCombinedInferenceAndNonInferenceQuery(queryBuilder, indexInformation);
+
+        if (fieldWeight != null && fieldWeight.equals(1.0f) == false) {
+            inferenceQuery.boost(fieldWeight);
+        }
+
+        return inferenceQuery;
+    }
+
     @Override
     public String getQueryName() {
         return MatchQueryBuilder.NAME;

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

@@ -8,6 +8,7 @@
 package org.elasticsearch.xpack.inference.queries;
 
 import org.elasticsearch.index.query.BoolQueryBuilder;
+import org.elasticsearch.index.query.MatchQueryBuilder;
 import org.elasticsearch.index.query.MultiMatchQueryBuilder;
 import org.elasticsearch.index.query.QueryBuilder;
 
@@ -50,8 +51,7 @@ protected QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceI
         return inferenceQuery;
     }
 
-    @Override
-    protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation) {
+    private QueryBuilder buildCombinedInferenceAndNonInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation) {
         assert (queryBuilder instanceof MultiMatchQueryBuilder);
         MultiMatchQueryBuilder originalMultiMatchQueryBuilder = (MultiMatchQueryBuilder) queryBuilder;
 
@@ -82,6 +82,42 @@ protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery(QueryBuilder q
         return boolQueryBuilder;
     }
 
+    @Override
+    protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation, Float fieldWeight) {
+        assert (queryBuilder instanceof MultiMatchQueryBuilder);
+        MultiMatchQueryBuilder originalMultiMatchQueryBuilder = (MultiMatchQueryBuilder) queryBuilder;
+
+        MatchQueryBuilder matchQueryBuilder = new MatchQueryBuilder(indexInformation.fieldName(), getQuery(queryBuilder));
+
+        // Create a copy for non-inference fields with only this specific field
+//        MultiMatchQueryBuilder multiMatchQueryBuilder = createSingleFieldMultiMatch(
+//            originalMultiMatchQueryBuilder,
+//            indexInformation.fieldName()
+//        );
+
+        BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder();
+
+        // Add semantic query for inference indices
+        boolQueryBuilder.should(
+            createSemanticSubQuery(
+                indexInformation.getInferenceIndices(),
+                indexInformation.fieldName(),
+                getQuery(matchQueryBuilder)
+            )
+        );
+
+        // Add regular query for non-inference indices
+        boolQueryBuilder.should(
+            createSubQueryForIndices(indexInformation.nonInferenceIndices(), matchQueryBuilder)
+        );
+
+        if (fieldWeight != null && fieldWeight.equals(1.0f) == false) {
+            boolQueryBuilder.boost(fieldWeight);
+        }
+        boolQueryBuilder.queryName(queryBuilder.queryName());
+        return boolQueryBuilder;
+    }
+
     @Override
     public String getQueryName() {
         return MultiMatchQueryBuilder.NAME;

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

@@ -53,7 +53,7 @@ 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.
-                finalQueryBuilder.should(buildCombinedInferenceAndNonInferenceQuery(queryBuilder, indexInformation));
+                finalQueryBuilder.should(buildCombinedInferenceAndNonInferenceQuery(queryBuilder, indexInformation, fieldWeight));
             } 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.
@@ -93,9 +93,15 @@ public QueryBuilder interceptAndRewrite(QueryRewriteContext context, QueryBuilde
      * @param indexInformation {@link InferenceIndexInformationForField}
      * @return {@link QueryBuilder}
      */
+//    protected abstract QueryBuilder buildCombinedInferenceAndNonInferenceQuery(
+//        QueryBuilder queryBuilder,
+//        InferenceIndexInformationForField indexInformation
+//    );
+
     protected abstract QueryBuilder buildCombinedInferenceAndNonInferenceQuery(
         QueryBuilder queryBuilder,
-        InferenceIndexInformationForField indexInformation
+        InferenceIndexInformationForField indexInformation,
+        Float fieldWeight
     );
 
     private InferenceIndexInformationForField resolveIndicesForField(String fieldName, ResolvedIndices resolvedIndices) {

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

@@ -83,7 +83,6 @@ private QueryBuilder buildInferenceQueryWithMultipleInferenceIds(
         return boolQueryBuilder;
     }
 
-    @Override
     protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery(
         QueryBuilder queryBuilder,
         InferenceIndexInformationForField indexInformation
@@ -121,6 +120,20 @@ protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery(
         return boolQueryBuilder;
     }
 
+    protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery(
+        QueryBuilder queryBuilder,
+        InferenceIndexInformationForField indexInformation,
+        Float fieldWeight
+    ) {
+        QueryBuilder inferenceQuery = buildCombinedInferenceAndNonInferenceQuery(queryBuilder, indexInformation);
+
+        if (fieldWeight != null && fieldWeight.equals(1.0f) == false) {
+            inferenceQuery.boost(fieldWeight);
+        }
+
+        return inferenceQuery;
+    }
+
     private QueryBuilder buildNestedQueryFromSparseVectorQuery(QueryBuilder queryBuilder, String searchInferenceId) {
         assert (queryBuilder instanceof SparseVectorQueryBuilder);
         SparseVectorQueryBuilder sparseVectorQueryBuilder = (SparseVectorQueryBuilder) queryBuilder;