Fix Semantic Query Rewrite Interception Drops Boosts

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

@@ -54,7 +54,7 @@ public void cleanup() {
     }
 
     public void testKnnQueryWithVectorBuilderIsInterceptedAndRewritten() throws IOException {
-        float boost = randomFloat() * 5;
+        float boost = randomFloatBetween(1, 10, true);
         String queryName = randomAlphaOfLength(5);
         Map<String, InferenceFieldMetadata> inferenceFields = Map.of(
             FIELD_NAME,
@@ -69,13 +69,17 @@ public void testKnnQueryWithVectorBuilderIsInterceptedAndRewritten() throws IOEx
     }
 
     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);
     }
 

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

@@ -94,7 +94,7 @@ public void testBoostAndQueryNameInMatchQueryRewrite() throws IOException {
             rewritten instanceof InterceptedQueryBuilderWrapper
         );
         InterceptedQueryBuilderWrapper intercepted = (InterceptedQueryBuilderWrapper) rewritten;
-        assertEquals(5.0, intercepted.boost(), 0.0f);
+        assertEquals(BOOST, intercepted.boost(), 0.0f);
         assertEquals(QUERY_NAME, intercepted.queryName());
         assertTrue(intercepted.queryBuilder instanceof SemanticQueryBuilder);
         SemanticQueryBuilder semanticQueryBuilder = (SemanticQueryBuilder) intercepted.queryBuilder;

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

@@ -52,7 +52,7 @@ public void cleanup() {
     }
 
     public void testSparseVectorQueryOnInferenceFieldIsInterceptedAndRewritten() throws IOException {
-        float boost = randomFloat() * 5;
+        float boost = randomFloatBetween(1, 10, true);
         String queryName = randomAlphaOfLength(5);
         Map<String, InferenceFieldMetadata> inferenceFields = Map.of(
             FIELD_NAME,
@@ -66,12 +66,16 @@ public void testSparseVectorQueryOnInferenceFieldIsInterceptedAndRewritten() thr
     }
 
     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);
     }
 

x-pack/plugin/inference/src/yamlRestTest/resources/rest-api-spec/test/inference/45_semantic_text_match.yml

@@ -322,7 +322,7 @@ setup:
 
   - match: { hits.total.value: 1 }
   - match: { hits.hits.0._id: "doc_1" }
-  - close_to: { hits.hits.0._score: { value: 5.700229E18, error: 1e15 } }
+  - close_to: { hits.hits.0._score: { value: 5.700229E18, error: 1e-7 } }
   - not_exists: hits.hits.0.matched_queries
 
   - do:
@@ -341,7 +341,7 @@ setup:
 
   - match: { hits.total.value: 1 }
   - match: { hits.hits.0._id: "doc_1" }
-  - close_to: { hits.hits.0._score: { value: 2.8501142E19, error: 1e15 } }
+  - close_to: { hits.hits.0._score: { value: 2.8501142E19, error: 1e-7 } }
   - match: { hits.hits.0.matched_queries: [ "i-like-naming-my-queries" ] }
 
 ---
@@ -406,7 +406,7 @@ setup:
   - match: { hits.total.value: 2 }
   - match: { hits.hits.0._id: "doc_1" }
   - match: { hits.hits.1._id: "doc_2" }
-  - close_to: { hits.hits.0._score: { value: 1.1140361E19, error: 1e15 } }
+  - close_to: { hits.hits.0._score: { value: 1.1140361E19, error: 1e-7 } }
   - not_exists: hits.hits.0.matched_queries
   - close_to: { hits.hits.1._score: { value: 0.2876821, error: 1e-7 } }
   - not_exists: hits.hits.1.matched_queries
@@ -428,7 +428,7 @@ setup:
   - match: { hits.total.value: 2 }
   - match: { hits.hits.0._id: "doc_1" }
   - match: { hits.hits.1._id: "doc_2" }
-  - close_to: { hits.hits.0._score: { value: 5.5701804E19, error: 1e15 } }
+  - close_to: { hits.hits.0._score: { value: 5.5701804E19, error: 1e-7 } }
   - match: { hits.hits.0.matched_queries: [ "i-like-naming-my-queries" ] }
   - close_to: { hits.hits.1._score: { value: 1.4384103, error: 1e-7 } }
   - match: { hits.hits.1.matched_queries: [ "i-like-naming-my-queries" ] }

x-pack/plugin/inference/src/yamlRestTest/resources/rest-api-spec/test/inference/46_semantic_text_sparse_vector.yml

@@ -291,7 +291,7 @@ setup:
 
   - match: { hits.total.value: 1 }
   - match: { hits.hits.0._id: "doc_1" }
-  - close_to: { hits.hits.0._score: { value: 5.700229E18, error: 1e15 } }
+  - close_to: { hits.hits.0._score: { value: 5.700229E18, error: 1e-7 } }
   - not_exists: hits.hits.0.matched_queries
 
   - do:
@@ -310,7 +310,7 @@ setup:
 
   - match: { hits.total.value: 1 }
   - match: { hits.hits.0._id: "doc_1" }
-  - close_to: { hits.hits.0._score: { value: 2.8501142E19, error: 1e15 } }
+  - close_to: { hits.hits.0._score: { value: 2.8501142E19, error: 1e-7 } }
   - match: { hits.hits.0.matched_queries: [ "i-like-naming-my-queries" ] }
 
 ---
@@ -376,7 +376,7 @@ setup:
   - match: { hits.total.value: 2 }
   - match: { hits.hits.0._id: "doc_1" }
   - match: { hits.hits.1._id: "doc_2" }
-  - close_to: { hits.hits.0._score: { value: 5.700229E18, error: 1e15 } }
+  - close_to: { hits.hits.0._score: { value: 5.700229E18, error: 1e-7 } }
   - not_exists: hits.hits.0.matched_queries
   - close_to: { hits.hits.1._score: { value: 1.3455845E10, error: 1e-7 } }
   - not_exists: hits.hits.1.matched_queries
@@ -399,7 +399,7 @@ setup:
   - match: { hits.total.value: 2 }
   - match: { hits.hits.0._id: "doc_1" }
   - match: { hits.hits.1._id: "doc_2" }
-  - close_to: { hits.hits.0._score: { value: 2.8501142E19, error: 1e15 } }
+  - close_to: { hits.hits.0._score: { value: 2.8501142E19, error: 1e-7 } }
   - match: { hits.hits.0.matched_queries: [ "i-like-naming-my-queries" ] }
   - close_to: { hits.hits.1._score: { value: 6.7279225E10, error: 1e-7 } }
   - match: { hits.hits.1.matched_queries: [ "i-like-naming-my-queries" ] }

x-pack/plugin/inference/src/yamlRestTest/resources/rest-api-spec/test/inference/47_semantic_text_knn.yml

@@ -452,7 +452,7 @@ setup:
 
   - match: { hits.total.value: 1 }
   - match: { hits.hits.0._id: "doc_1" }
-  - close_to: { hits.hits.0._score: { value: 0.9984111, error: 1e15 } }
+  - close_to: { hits.hits.0._score: { value: 0.9981499, error: 1e-7 } }
   - not_exists: hits.hits.0.matched_queries
 
   - do: