additional tests and refactoring

Author: markjhoy

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/search/SparseVectorQueryBuilder.java

@@ -42,6 +42,7 @@
 import java.util.Map;
 import java.util.Objects;
 
+import static org.elasticsearch.TransportVersions.SPARSE_VECTOR_FIELD_PRUNING_OPTIONS_8_19;
 import static org.elasticsearch.xcontent.ConstructingObjectParser.constructorArg;
 import static org.elasticsearch.xcontent.ConstructingObjectParser.optionalConstructorArg;
 import static org.elasticsearch.xpack.core.ClientHelper.ML_ORIGIN;
@@ -124,7 +125,9 @@ public SparseVectorQueryBuilder(
     public SparseVectorQueryBuilder(StreamInput in) throws IOException {
         super(in);
         this.fieldName = in.readString();
-        if (in.getTransportVersion().onOrAfter(TransportVersions.SPARSE_VECTOR_FIELD_PRUNING_OPTIONS)) {
+        if (in.getTransportVersion().isPatchFrom(SPARSE_VECTOR_FIELD_PRUNING_OPTIONS_8_19) ||
+            in.getTransportVersion().onOrAfter(TransportVersions.SPARSE_VECTOR_FIELD_PRUNING_OPTIONS)
+        ) {
             this.shouldPruneTokens = in.readOptionalBoolean();
         } else {
             this.shouldPruneTokens = in.readBoolean();
@@ -177,7 +180,9 @@ protected void doWriteTo(StreamOutput out) throws IOException {
         }
 
         out.writeString(fieldName);
-        if (out.getTransportVersion().onOrAfter(TransportVersions.SPARSE_VECTOR_FIELD_PRUNING_OPTIONS)) {
+        if (out.getTransportVersion().isPatchFrom(SPARSE_VECTOR_FIELD_PRUNING_OPTIONS_8_19)
+            || out.getTransportVersion().onOrAfter(TransportVersions.SPARSE_VECTOR_FIELD_PRUNING_OPTIONS)
+        ) {
             out.writeOptionalBoolean(shouldPruneTokens);
         } else {
             out.writeBoolean(shouldPruneTokens != null && shouldPruneTokens);

x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/search/SparseVectorQueryBuilderTests.java

@@ -155,19 +155,26 @@ protected Object simulateMethod(Method method, Object[] args) {
 
     @Override
     protected void initializeAdditionalMappings(MapperService mapperService) throws IOException {
-        if (shouldInjectSparseVectorIndexOptions()) {
-            addSparseVectorIndexOptionsMapping(mapperService);
-            return;
-        }
-
         mapperService.merge(
             "_doc",
-            new CompressedXContent(Strings.toString(PutMappingRequest.simpleMapping(SPARSE_VECTOR_FIELD, "type=sparse_vector"))),
+            new CompressedXContent(getTestSparseVectorIndexMapping()),
             MapperService.MergeReason.MAPPING_UPDATE
         );
     }
 
-    private boolean shouldInjectSparseVectorIndexOptions() {
+    private String getTestSparseVectorIndexMapping() {
+        if (currentTestHasIndexOptions()) {
+            return "{\"properties\":{\""
+                + SPARSE_VECTOR_FIELD
+                + "\":{\"type\":\"sparse_vector\",\"index_options\""
+                + ":{\"prune\":true,\"pruning_config\":{\"tokens_freq_ratio_threshold\""
+                + ":12,\"tokens_weight_threshold\":0.6}}}}}";
+        }
+
+        return Strings.toString(PutMappingRequest.simpleMapping(SPARSE_VECTOR_FIELD, "type=sparse_vector"));
+    }
+
+    private boolean currentTestHasIndexOptions() {
         Class<?> clazz = this.getClass();
         Class<InjectSparseVectorIndexOptions> injectSparseVectorIndexOptions = InjectSparseVectorIndexOptions.class;
 
@@ -179,15 +186,6 @@ private boolean shouldInjectSparseVectorIndexOptions() {
         }
     }
 
-    private void addSparseVectorIndexOptionsMapping(MapperService mapperService) throws IOException {
-        String addIndexOptionsTemplate = "{\"properties\":{\""
-            + SPARSE_VECTOR_FIELD
-            + "\":{\"type\":\"sparse_vector\",\"index_options\""
-            + ":{\"prune\":true,\"pruning_config\":{\"tokens_freq_ratio_threshold\""
-            + ":12,\"tokens_weight_threshold\":0.6}}}}}";
-        mapperService.merge("_doc", new CompressedXContent(addIndexOptionsTemplate), MapperService.MergeReason.MAPPING_UPDATE);
-    }
-
     @Override
     protected void doAssertLuceneQuery(SparseVectorQueryBuilder queryBuilder, Query query, SearchExecutionContext context) {
         assertThat(query, instanceOf(SparseVectorQueryWrapper.class));
@@ -283,7 +281,7 @@ private void testDoToQuery(SparseVectorQueryBuilder queryBuilder, SearchExecutio
 
         assertTrue(query instanceof SparseVectorQueryWrapper);
         var sparseQuery = (SparseVectorQueryWrapper) query;
-        if (queryBuilder.shouldPruneTokens()) {
+        if (queryBuilder.shouldPruneTokens() || currentTestHasIndexOptions()) {
             // It's possible that all documents were pruned for aggressive pruning configurations
             assertTrue(sparseQuery.getTermsQuery() instanceof BooleanQuery || sparseQuery.getTermsQuery() instanceof MatchNoDocsQuery);
         } else {
@@ -386,31 +384,33 @@ public void testThatWeCorrectlyRewriteQueryIntoVectors() {
     public void testItUsesIndexOptionsDefaults() throws IOException {
         withSearchIndex((context) -> {
             try {
-                SparseVectorQueryBuilder builder = new SparseVectorQueryBuilder(
-                    SPARSE_VECTOR_FIELD,
-                    WEIGHTED_TOKENS,
-                    null,
-                    null,
-                    null,
-                    null
-                );
-                Query query = builder.doToQuery(context);
-
+                SparseVectorQueryBuilder builder = createTestQueryBuilder(null);
+                assertFalse(builder.shouldPruneTokens());
+                testDoToQuery(builder, context);
             } catch (IOException ex) {
                 throw new RuntimeException(ex);
             }
         });
     }
 
     @InjectSparseVectorIndexOptions
-    public void testItOverridesIndexOptionsDefaults() {
-
+    public void testItOverridesIndexOptionsDefaults() throws IOException {
+        withSearchIndex((context) -> {
+            try {
+                TokenPruningConfig pruningConfig = new TokenPruningConfig(2, 0.3f, false);
+                SparseVectorQueryBuilder builder = createTestQueryBuilder(pruningConfig);
+                assertTrue(builder.shouldPruneTokens());
+                testDoToQuery(builder, context);
+            } catch (IOException ex) {
+                throw new RuntimeException(ex);
+            }
+        });
     }
 
     @InjectSparseVectorIndexOptions
-    public void testToQueryWithIndexOptions() throws IOException {
+    public void testToQueryRewriteWithIndexOptions() throws IOException {
         withSearchIndex((context) -> {
-            SparseVectorQueryBuilder queryBuilder = createTestQueryBuilder();
+            SparseVectorQueryBuilder queryBuilder = createTestQueryBuilder(null);
             try {
                 if (queryBuilder.getQueryVectors() == null) {
                     QueryBuilder rewrittenQueryBuilder = rewriteAndFetch(queryBuilder, context);
@@ -424,15 +424,4 @@ public void testToQueryWithIndexOptions() throws IOException {
             }
         });
     }
-
-    @InjectSparseVectorIndexOptions
-    public void testWeCorrectlyRewriteQueryIntoVectorsWithIndexOptions() {
-        SearchExecutionContext searchExecutionContext = createSearchExecutionContext();
-
-        SparseVectorQueryBuilder queryBuilder = createTestQueryBuilder(null);
-        QueryBuilder rewrittenQueryBuilder = rewriteAndFetch(queryBuilder, searchExecutionContext);
-        assertTrue(rewrittenQueryBuilder instanceof SparseVectorQueryBuilder);
-        assertEquals(queryBuilder.shouldPruneTokens(), ((SparseVectorQueryBuilder) rewrittenQueryBuilder).shouldPruneTokens());
-        assertNotNull(((SparseVectorQueryBuilder) rewrittenQueryBuilder).getQueryVectors());
-    }
 }