Fix Semantic Query Rewrite Interception Drops Boosts

docs/changelog/129282.yaml

@@ -0,0 +1,6 @@
+pr: 129282
+summary: "Fix query rewrite logic to preserve `boosts` and `queryName` for `match`,\
+  \ `knn`, and `sparse_vector` queries on semantic_text fields"
+area: Search
+type: bug
+issues: []

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

@@ -37,6 +37,9 @@ public class InferenceFeatures implements FeatureSpecification {
     private static final NodeFeature TEST_RULE_RETRIEVER_WITH_INDICES_THAT_DONT_RETURN_RANK_DOCS = new NodeFeature(
         "test_rule_retriever.with_indices_that_dont_return_rank_docs"
     );
+    private static final NodeFeature SEMANTIC_QUERY_REWRITE_INTERCEPTORS_PROPAGATE_BOOST_AND_QUERY_NAME_FIX = new NodeFeature(
+        "semantic_query_rewrite_interceptors.propagate_boost_and_query_name_fix"
+    );
     private static final NodeFeature SEMANTIC_TEXT_MATCH_ALL_HIGHLIGHTER = new NodeFeature("semantic_text.match_all_highlighter");
     private static final NodeFeature COHERE_V2_API = new NodeFeature("inference.cohere.v2");
 
@@ -68,7 +71,8 @@ public Set<NodeFeature> getTestFeatures() {
             SEMANTIC_TEXT_EXCLUDE_SUB_FIELDS_FROM_FIELD_CAPS,
             SEMANTIC_TEXT_INDEX_OPTIONS,
             COHERE_V2_API,
-            SEMANTIC_TEXT_INDEX_OPTIONS_WITH_DEFAULTS
+            SEMANTIC_TEXT_INDEX_OPTIONS_WITH_DEFAULTS,
+            SEMANTIC_QUERY_REWRITE_INTERCEPTORS_PROPAGATE_BOOST_AND_QUERY_NAME_FIX
         );
     }
 }

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

@@ -52,16 +52,20 @@ protected QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceI
         assert (queryBuilder instanceof KnnVectorQueryBuilder);
         KnnVectorQueryBuilder knnVectorQueryBuilder = (KnnVectorQueryBuilder) queryBuilder;
         Map<String, List<String>> inferenceIdsIndices = indexInformation.getInferenceIdsIndices();
+        QueryBuilder finalQueryBuilder;
         if (inferenceIdsIndices.size() == 1) {
             // Simple case, everything uses the same inference ID
             Map.Entry<String, List<String>> inferenceIdIndex = inferenceIdsIndices.entrySet().iterator().next();
             String searchInferenceId = inferenceIdIndex.getKey();
             List<String> indices = inferenceIdIndex.getValue();
-            return buildNestedQueryFromKnnVectorQuery(knnVectorQueryBuilder, indices, searchInferenceId);
+            finalQueryBuilder = buildNestedQueryFromKnnVectorQuery(knnVectorQueryBuilder, indices, searchInferenceId);
         } else {
             // Multiple inference IDs, construct a boolean query
-            return buildInferenceQueryWithMultipleInferenceIds(knnVectorQueryBuilder, inferenceIdsIndices);
+            finalQueryBuilder = buildInferenceQueryWithMultipleInferenceIds(knnVectorQueryBuilder, inferenceIdsIndices);
         }
+        finalQueryBuilder.boost(queryBuilder.boost());
+        finalQueryBuilder.queryName(queryBuilder.queryName());
+        return finalQueryBuilder;
     }
 
     private QueryBuilder buildInferenceQueryWithMultipleInferenceIds(
@@ -102,6 +106,8 @@ protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery(
                 )
             );
         }
+        boolQueryBuilder.boost(queryBuilder.boost());
+        boolQueryBuilder.queryName(queryBuilder.queryName());
         return boolQueryBuilder;
     }
 

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

@@ -36,7 +36,10 @@ protected String getQuery(QueryBuilder queryBuilder) {
 
     @Override
     protected QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation) {
-        return new SemanticQueryBuilder(indexInformation.fieldName(), getQuery(queryBuilder), false);
+        SemanticQueryBuilder semanticQueryBuilder = new SemanticQueryBuilder(indexInformation.fieldName(), getQuery(queryBuilder), false);
+        semanticQueryBuilder.boost(queryBuilder.boost());
+        semanticQueryBuilder.queryName(queryBuilder.queryName());
+        return semanticQueryBuilder;
     }
 
     @Override
@@ -45,7 +48,10 @@ protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery(
         InferenceIndexInformationForField indexInformation
     ) {
         assert (queryBuilder instanceof MatchQueryBuilder);
-        MatchQueryBuilder matchQueryBuilder = (MatchQueryBuilder) queryBuilder;
+        MatchQueryBuilder originalMatchQueryBuilder = (MatchQueryBuilder) queryBuilder;
+        // Create a copy for non-inference fields without boost and _name
+        MatchQueryBuilder matchQueryBuilder = copyMatchQueryBuilder(originalMatchQueryBuilder);
+
         BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder();
         boolQueryBuilder.should(
             createSemanticSubQuery(
@@ -55,11 +61,33 @@ protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery(
             )
         );
         boolQueryBuilder.should(createSubQueryForIndices(indexInformation.nonInferenceIndices(), matchQueryBuilder));
+        boolQueryBuilder.boost(queryBuilder.boost());
+        boolQueryBuilder.queryName(queryBuilder.queryName());
         return boolQueryBuilder;
     }
 
     @Override
     public String getQueryName() {
         return MatchQueryBuilder.NAME;
     }
+
+    private MatchQueryBuilder copyMatchQueryBuilder(MatchQueryBuilder queryBuilder) {
+        MatchQueryBuilder matchQueryBuilder = new MatchQueryBuilder(queryBuilder.fieldName(), queryBuilder.value());
+        matchQueryBuilder.operator(queryBuilder.operator());
+        matchQueryBuilder.prefixLength(queryBuilder.prefixLength());
+        matchQueryBuilder.maxExpansions(queryBuilder.maxExpansions());
+        matchQueryBuilder.fuzzyTranspositions(queryBuilder.fuzzyTranspositions());
+        matchQueryBuilder.lenient(queryBuilder.lenient());
+        matchQueryBuilder.zeroTermsQuery(queryBuilder.zeroTermsQuery());
+        matchQueryBuilder.analyzer(queryBuilder.analyzer());
+        matchQueryBuilder.minimumShouldMatch(queryBuilder.minimumShouldMatch());
+        matchQueryBuilder.fuzzyRewrite(queryBuilder.fuzzyRewrite());
+
+        if (queryBuilder.fuzziness() != null) {
+            matchQueryBuilder.fuzziness(queryBuilder.fuzziness());
+        }
+
+        matchQueryBuilder.autoGenerateSynonymsPhraseQuery(queryBuilder.autoGenerateSynonymsPhraseQuery());
+        return matchQueryBuilder;
+    }
 }

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

@@ -43,14 +43,18 @@ protected String getQuery(QueryBuilder queryBuilder) {
     @Override
     protected QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation) {
         Map<String, List<String>> inferenceIdsIndices = indexInformation.getInferenceIdsIndices();
+        QueryBuilder finalQueryBuilder;
         if (inferenceIdsIndices.size() == 1) {
             // Simple case, everything uses the same inference ID
             String searchInferenceId = inferenceIdsIndices.keySet().iterator().next();
-            return buildNestedQueryFromSparseVectorQuery(queryBuilder, searchInferenceId);
+            finalQueryBuilder = buildNestedQueryFromSparseVectorQuery(queryBuilder, searchInferenceId);
         } else {
             // Multiple inference IDs, construct a boolean query
-            return buildInferenceQueryWithMultipleInferenceIds(queryBuilder, inferenceIdsIndices);
+            finalQueryBuilder = buildInferenceQueryWithMultipleInferenceIds(queryBuilder, inferenceIdsIndices);
         }
+        finalQueryBuilder.queryName(queryBuilder.queryName());
+        finalQueryBuilder.boost(queryBuilder.boost());
+        return finalQueryBuilder;
     }
 
     private QueryBuilder buildInferenceQueryWithMultipleInferenceIds(
@@ -79,7 +83,19 @@ protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery(
         Map<String, List<String>> inferenceIdsIndices = indexInformation.getInferenceIdsIndices();
 
         BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder();
-        boolQueryBuilder.should(createSubQueryForIndices(indexInformation.nonInferenceIndices(), sparseVectorQueryBuilder));
+        boolQueryBuilder.should(
+            createSubQueryForIndices(
+                indexInformation.nonInferenceIndices(),
+                new SparseVectorQueryBuilder(
+                    sparseVectorQueryBuilder.getFieldName(),
+                    sparseVectorQueryBuilder.getQueryVectors(),
+                    sparseVectorQueryBuilder.getInferenceId(),
+                    sparseVectorQueryBuilder.getQuery(),
+                    sparseVectorQueryBuilder.shouldPruneTokens(),
+                    sparseVectorQueryBuilder.getTokenPruningConfig()
+                )
+            )
+        );
         // We always perform nested subqueries on semantic_text fields, to support
         // sparse_vector queries using query vectors.
         for (String inferenceId : inferenceIdsIndices.keySet()) {
@@ -90,6 +106,8 @@ protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery(
                 )
             );
         }
+        boolQueryBuilder.boost(queryBuilder.boost());
+        boolQueryBuilder.queryName(queryBuilder.queryName());
         return boolQueryBuilder;
     }
 

x-pack/plugin/inference/src/test/java/org/elasticsearch/index/query/SemanticKnnVectorQueryRewriteInterceptorTests.java

@@ -54,24 +54,32 @@ public void cleanup() {
     }
 
     public void testKnnQueryWithVectorBuilderIsInterceptedAndRewritten() throws IOException {
+        float boost = randomFloatBetween(1, 10, true);
+        String queryName = randomAlphaOfLength(5);
         Map<String, InferenceFieldMetadata> inferenceFields = Map.of(
             FIELD_NAME,
             new InferenceFieldMetadata(index.getName(), INFERENCE_ID, new String[] { FIELD_NAME }, null)
         );
         QueryRewriteContext context = createQueryRewriteContext(inferenceFields);
         QueryVectorBuilder queryVectorBuilder = new TextEmbeddingQueryVectorBuilder(INFERENCE_ID, QUERY);
         KnnVectorQueryBuilder original = new KnnVectorQueryBuilder(FIELD_NAME, queryVectorBuilder, 10, 100, null);
+        original.boost(boost);
+        original.queryName(queryName);
         testRewrittenInferenceQuery(context, original);
     }
 
     public void testKnnWithQueryBuilderWithoutInferenceIdIsInterceptedAndRewritten() throws IOException {
+        float boost = randomFloatBetween(1, 10, true);
+        String queryName = randomAlphaOfLength(5);
         Map<String, InferenceFieldMetadata> inferenceFields = Map.of(
             FIELD_NAME,
             new InferenceFieldMetadata(index.getName(), INFERENCE_ID, new String[] { FIELD_NAME }, null)
         );
         QueryRewriteContext context = createQueryRewriteContext(inferenceFields);
         QueryVectorBuilder queryVectorBuilder = new TextEmbeddingQueryVectorBuilder(null, QUERY);
         KnnVectorQueryBuilder original = new KnnVectorQueryBuilder(FIELD_NAME, queryVectorBuilder, 10, 100, null);
+        original.boost(boost);
+        original.queryName(queryName);
         testRewrittenInferenceQuery(context, original);
     }
 
@@ -82,14 +90,23 @@ private void testRewrittenInferenceQuery(QueryRewriteContext context, KnnVectorQ
             rewritten instanceof InterceptedQueryBuilderWrapper
         );
         InterceptedQueryBuilderWrapper intercepted = (InterceptedQueryBuilderWrapper) rewritten;
+        assertEquals(original.boost(), intercepted.boost(), 0.0f);
+        assertEquals(original.queryName(), intercepted.queryName());
         assertTrue(intercepted.queryBuilder instanceof NestedQueryBuilder);
+
         NestedQueryBuilder nestedQueryBuilder = (NestedQueryBuilder) intercepted.queryBuilder;
+        assertEquals(original.boost(), nestedQueryBuilder.boost(), 0.0f);
+        assertEquals(original.queryName(), nestedQueryBuilder.queryName());
         assertEquals(SemanticTextField.getChunksFieldName(FIELD_NAME), nestedQueryBuilder.path());
+
         QueryBuilder innerQuery = nestedQueryBuilder.query();
         assertTrue(innerQuery instanceof KnnVectorQueryBuilder);
         KnnVectorQueryBuilder knnVectorQueryBuilder = (KnnVectorQueryBuilder) innerQuery;
+        assertEquals(1.0f, knnVectorQueryBuilder.boost(), 0.0f);
+        assertNull(knnVectorQueryBuilder.queryName());
         assertEquals(SemanticTextField.getEmbeddingsFieldName(FIELD_NAME), knnVectorQueryBuilder.getFieldName());
         assertTrue(knnVectorQueryBuilder.queryVectorBuilder() instanceof TextEmbeddingQueryVectorBuilder);
+
         TextEmbeddingQueryVectorBuilder textEmbeddingQueryVectorBuilder = (TextEmbeddingQueryVectorBuilder) knnVectorQueryBuilder
             .queryVectorBuilder();
         assertEquals(QUERY, textEmbeddingQueryVectorBuilder.getModelText());

x-pack/plugin/inference/src/test/java/org/elasticsearch/index/query/SemanticMatchQueryRewriteInterceptorTests.java

@@ -36,6 +36,8 @@ public class SemanticMatchQueryRewriteInterceptorTests extends ESTestCase {
 
     private static final String FIELD_NAME = "fieldName";
     private static final String VALUE = "value";
+    private static final String QUERY_NAME = "match_query";
+    private static final float BOOST = 5.0f;
 
     @Before
     public void setup() {
@@ -79,10 +81,38 @@ public void testMatchQueryOnNonInferenceFieldRemainsMatchQuery() throws IOExcept
         assertEquals(original, rewritten);
     }
 
+    public void testBoostAndQueryNameInMatchQueryRewrite() throws IOException {
+        Map<String, InferenceFieldMetadata> inferenceFields = Map.of(
+            FIELD_NAME,
+            new InferenceFieldMetadata(index.getName(), "inferenceId", new String[] { FIELD_NAME }, null)
+        );
+        QueryRewriteContext context = createQueryRewriteContext(inferenceFields);
+        QueryBuilder original = createTestQueryBuilderWithBoostAndQueryName();
+        QueryBuilder rewritten = original.rewrite(context);
+        assertTrue(
+            "Expected query to be intercepted, but was [" + rewritten.getClass().getName() + "]",
+            rewritten instanceof InterceptedQueryBuilderWrapper
+        );
+        InterceptedQueryBuilderWrapper intercepted = (InterceptedQueryBuilderWrapper) rewritten;
+        assertEquals(BOOST, intercepted.boost(), 0.0f);
+        assertEquals(QUERY_NAME, intercepted.queryName());
+        assertTrue(intercepted.queryBuilder instanceof SemanticQueryBuilder);
+        SemanticQueryBuilder semanticQueryBuilder = (SemanticQueryBuilder) intercepted.queryBuilder;
+        assertEquals(FIELD_NAME, semanticQueryBuilder.getFieldName());
+        assertEquals(VALUE, semanticQueryBuilder.getQuery());
+    }
+
     private MatchQueryBuilder createTestQueryBuilder() {
         return new MatchQueryBuilder(FIELD_NAME, VALUE);
     }
 
+    private MatchQueryBuilder createTestQueryBuilderWithBoostAndQueryName() {
+        MatchQueryBuilder queryBuilder = createTestQueryBuilder();
+        queryBuilder.boost(BOOST);
+        queryBuilder.queryName(QUERY_NAME);
+        return queryBuilder;
+    }
+
     private QueryRewriteContext createQueryRewriteContext(Map<String, InferenceFieldMetadata> inferenceFields) {
         IndexMetadata indexMetadata = IndexMetadata.builder(index.getName())
             .settings(

x-pack/plugin/inference/src/test/java/org/elasticsearch/index/query/SemanticSparseVectorQueryRewriteInterceptorTests.java

@@ -52,62 +52,68 @@ public void cleanup() {
     }
 
     public void testSparseVectorQueryOnInferenceFieldIsInterceptedAndRewritten() throws IOException {
+        float boost = randomFloatBetween(1, 10, true);
+        String queryName = randomAlphaOfLength(5);
         Map<String, InferenceFieldMetadata> inferenceFields = Map.of(
             FIELD_NAME,
             new InferenceFieldMetadata(index.getName(), "inferenceId", new String[] { FIELD_NAME }, null)
         );
         QueryRewriteContext context = createQueryRewriteContext(inferenceFields);
         QueryBuilder original = new SparseVectorQueryBuilder(FIELD_NAME, INFERENCE_ID, QUERY);
-        QueryBuilder rewritten = original.rewrite(context);
-        assertTrue(
-            "Expected query to be intercepted, but was [" + rewritten.getClass().getName() + "]",
-            rewritten instanceof InterceptedQueryBuilderWrapper
-        );
-        InterceptedQueryBuilderWrapper intercepted = (InterceptedQueryBuilderWrapper) rewritten;
-        assertTrue(intercepted.queryBuilder instanceof NestedQueryBuilder);
-        NestedQueryBuilder nestedQueryBuilder = (NestedQueryBuilder) intercepted.queryBuilder;
-        assertEquals(SemanticTextField.getChunksFieldName(FIELD_NAME), nestedQueryBuilder.path());
-        QueryBuilder innerQuery = nestedQueryBuilder.query();
-        assertTrue(innerQuery instanceof SparseVectorQueryBuilder);
-        SparseVectorQueryBuilder sparseVectorQueryBuilder = (SparseVectorQueryBuilder) innerQuery;
-        assertEquals(SemanticTextField.getEmbeddingsFieldName(FIELD_NAME), sparseVectorQueryBuilder.getFieldName());
-        assertEquals(INFERENCE_ID, sparseVectorQueryBuilder.getInferenceId());
-        assertEquals(QUERY, sparseVectorQueryBuilder.getQuery());
+        original.boost(boost);
+        original.queryName(queryName);
+        testRewrittenInferenceQuery(context, original);
     }
 
     public void testSparseVectorQueryOnInferenceFieldWithoutInferenceIdIsInterceptedAndRewritten() throws IOException {
+        float boost = randomFloatBetween(1, 10, true);
+        String queryName = randomAlphaOfLength(5);
         Map<String, InferenceFieldMetadata> inferenceFields = Map.of(
             FIELD_NAME,
             new InferenceFieldMetadata(index.getName(), "inferenceId", new String[] { FIELD_NAME }, null)
         );
         QueryRewriteContext context = createQueryRewriteContext(inferenceFields);
         QueryBuilder original = new SparseVectorQueryBuilder(FIELD_NAME, null, QUERY);
+        original.boost(boost);
+        original.queryName(queryName);
+        testRewrittenInferenceQuery(context, original);
+    }
+
+    public void testSparseVectorQueryOnNonInferenceFieldRemainsUnchanged() throws IOException {
+        QueryRewriteContext context = createQueryRewriteContext(Map.of()); // No inference fields
+        QueryBuilder original = new SparseVectorQueryBuilder(FIELD_NAME, INFERENCE_ID, QUERY);
+        QueryBuilder rewritten = original.rewrite(context);
+        assertTrue(
+            "Expected query to remain sparse_vector but was [" + rewritten.getClass().getName() + "]",
+            rewritten instanceof SparseVectorQueryBuilder
+        );
+        assertEquals(original, rewritten);
+    }
+
+    private void testRewrittenInferenceQuery(QueryRewriteContext context, QueryBuilder original) throws IOException {
         QueryBuilder rewritten = original.rewrite(context);
         assertTrue(
             "Expected query to be intercepted, but was [" + rewritten.getClass().getName() + "]",
             rewritten instanceof InterceptedQueryBuilderWrapper
         );
         InterceptedQueryBuilderWrapper intercepted = (InterceptedQueryBuilderWrapper) rewritten;
+        assertEquals(original.boost(), intercepted.boost(), 0.0f);
+        assertEquals(original.queryName(), intercepted.queryName());
+
         assertTrue(intercepted.queryBuilder instanceof NestedQueryBuilder);
         NestedQueryBuilder nestedQueryBuilder = (NestedQueryBuilder) intercepted.queryBuilder;
         assertEquals(SemanticTextField.getChunksFieldName(FIELD_NAME), nestedQueryBuilder.path());
+        assertEquals(original.boost(), nestedQueryBuilder.boost(), 0.0f);
+        assertEquals(original.queryName(), nestedQueryBuilder.queryName());
+
         QueryBuilder innerQuery = nestedQueryBuilder.query();
         assertTrue(innerQuery instanceof SparseVectorQueryBuilder);
         SparseVectorQueryBuilder sparseVectorQueryBuilder = (SparseVectorQueryBuilder) innerQuery;
         assertEquals(SemanticTextField.getEmbeddingsFieldName(FIELD_NAME), sparseVectorQueryBuilder.getFieldName());
         assertEquals(INFERENCE_ID, sparseVectorQueryBuilder.getInferenceId());
         assertEquals(QUERY, sparseVectorQueryBuilder.getQuery());
-    }
-
-    public void testSparseVectorQueryOnNonInferenceFieldRemainsUnchanged() throws IOException {
-        QueryRewriteContext context = createQueryRewriteContext(Map.of()); // No inference fields
-        QueryBuilder original = new SparseVectorQueryBuilder(FIELD_NAME, INFERENCE_ID, QUERY);
-        QueryBuilder rewritten = original.rewrite(context);
-        assertTrue(
-            "Expected query to remain sparse_vector but was [" + rewritten.getClass().getName() + "]",
-            rewritten instanceof SparseVectorQueryBuilder
-        );
-        assertEquals(original, rewritten);
+        assertEquals(1.0f, sparseVectorQueryBuilder.boost(), 0.0f);
+        assertNull(sparseVectorQueryBuilder.queryName());
     }
 
     private QueryRewriteContext createQueryRewriteContext(Map<String, InferenceFieldMetadata> inferenceFields) {