clean SparseVectorQueryBuilderTests

Author: markjhoy

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

@@ -44,10 +44,6 @@
 import org.elasticsearch.xpack.core.ml.inference.results.TextExpansionResults;
 
 import java.io.IOException;
-import java.lang.annotation.ElementType;
-import java.lang.annotation.Retention;
-import java.lang.annotation.RetentionPolicy;
-import java.lang.annotation.Target;
 import java.lang.reflect.Method;
 import java.util.ArrayList;
 import java.util.Collection;
@@ -65,17 +61,20 @@ public class SparseVectorQueryBuilderTests extends AbstractQueryTestCase<SparseV
     private static final int NUM_TOKENS = WEIGHTED_TOKENS.size();
     private final IndexVersion indexVersionToTest;
 
-    @Retention(RetentionPolicy.RUNTIME)
-    @Target(ElementType.METHOD)
-    public @interface InjectSparseVectorIndexOptions {
-    }
-
     public SparseVectorQueryBuilderTests() {
         // The sparse_vector field is not supported on versions 8.0 to 8.10. Because of this we'll only allow
         // index versions after its reintroduction.
         indexVersionToTest = randomBoolean()
-            ? IndexVersion.current()
-            : IndexVersionUtils.randomVersionBetween(random(), IndexVersions.NEW_SPARSE_VECTOR, IndexVersion.current());
+            ? IndexVersionUtils.randomVersionBetween(
+                random(),
+                IndexVersions.SPARSE_VECTOR_PRUNING_INDEX_OPTIONS_SUPPORT,
+                IndexVersion.current()
+            )
+            : IndexVersionUtils.randomVersionBetween(
+                random(),
+                IndexVersions.NEW_SPARSE_VECTOR,
+                IndexVersions.SPARSE_VECTOR_PRUNING_INDEX_OPTIONS_SUPPORT
+            );
     }
 
     @Override
@@ -159,50 +158,8 @@ protected Object simulateMethod(Method method, Object[] args) {
 
     @Override
     protected void initializeAdditionalMappings(MapperService mapperService) throws IOException {
-        mapperService.merge("_doc", new CompressedXContent(getTestSparseVectorIndexMapping()), MapperService.MergeReason.MAPPING_UPDATE);
-    }
-
-    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() {
-        if (indexVersionSupportsIndexOptions() == false) {
-            return false;
-        }
-
-        Class<?> clazz = this.getClass();
-        Class<InjectSparseVectorIndexOptions> injectSparseVectorIndexOptions = InjectSparseVectorIndexOptions.class;
-
-        try {
-            Method method = clazz.getMethod(this.getTestName());
-            return method.isAnnotationPresent(injectSparseVectorIndexOptions);
-        } catch (NoSuchMethodException e) {
-            return false;
-        }
-    }
-
-    private boolean indexVersionSupportsIndexOptions() {
-        if (indexVersionToTest.onOrAfter(IndexVersions.SPARSE_VECTOR_PRUNING_INDEX_OPTIONS_SUPPORT)) {
-            return true;
-        }
-
-        if (indexVersionToTest.between(
-            IndexVersions.SPARSE_VECTOR_PRUNING_INDEX_OPTIONS_SUPPORT_BACKPORT_8_X,
-            IndexVersions.UPGRADE_TO_LUCENE_10_0_0
-        )) {
-            return true;
-        }
-
-        return false;
+        String mapping = Strings.toString(PutMappingRequest.simpleMapping(SPARSE_VECTOR_FIELD, "type=sparse_vector"));
+        mapperService.merge("_doc", new CompressedXContent(mapping), MapperService.MergeReason.MAPPING_UPDATE);
     }
 
     @Override
@@ -242,15 +199,11 @@ private void withSearchIndex(CheckedConsumer<SearchExecutionContext, IOException
     @Override
     public void testCacheability() throws IOException {
         withSearchIndex((context) -> {
-            try {
-                SparseVectorQueryBuilder queryBuilder = createTestQueryBuilder();
-                QueryBuilder rewriteQuery = null;
-                rewriteQuery = rewriteQuery(queryBuilder, new SearchExecutionContext(context));
-                assertNotNull(rewriteQuery.toQuery(context));
-                assertTrue("query should be cacheable: " + queryBuilder.toString(), context.isCacheable());
-            } catch (IOException e) {
-                throw new RuntimeException(e);
-            }
+            SparseVectorQueryBuilder queryBuilder = createTestQueryBuilder();
+            QueryBuilder rewriteQuery = null;
+            rewriteQuery = rewriteQuery(queryBuilder, new SearchExecutionContext(context));
+            assertNotNull(rewriteQuery.toQuery(context));
+            assertTrue("query should be cacheable: " + queryBuilder.toString(), context.isCacheable());
         });
     }
 
@@ -260,12 +213,8 @@ public void testCacheability() throws IOException {
     @Override
     public void testMustRewrite() throws IOException {
         withSearchIndex((context) -> {
-            try {
-                SparseVectorQueryBuilder queryBuilder = createTestQueryBuilder();
-                queryBuilder.toQuery(context);
-            } catch (IOException e) {
-                throw new RuntimeException(e);
-            }
+            SparseVectorQueryBuilder queryBuilder = createTestQueryBuilder();
+            queryBuilder.toQuery(context);
         });
     }
 
@@ -275,17 +224,13 @@ public void testMustRewrite() throws IOException {
     @Override
     public void testToQuery() throws IOException {
         withSearchIndex((context) -> {
-            try {
-                SparseVectorQueryBuilder queryBuilder = createTestQueryBuilder();
-                if (queryBuilder.getQueryVectors() == null) {
-                    QueryBuilder rewrittenQueryBuilder = rewriteAndFetch(queryBuilder, context);
-                    assertTrue(rewrittenQueryBuilder instanceof SparseVectorQueryBuilder);
-                    testDoToQuery((SparseVectorQueryBuilder) rewrittenQueryBuilder, context);
-                } else {
-                    testDoToQuery(queryBuilder, context);
-                }
-            } catch (IOException e) {
-                throw new RuntimeException(e);
+            SparseVectorQueryBuilder queryBuilder = createTestQueryBuilder();
+            if (queryBuilder.getQueryVectors() == null) {
+                QueryBuilder rewrittenQueryBuilder = rewriteAndFetch(queryBuilder, context);
+                assertTrue(rewrittenQueryBuilder instanceof SparseVectorQueryBuilder);
+                testDoToQuery((SparseVectorQueryBuilder) rewrittenQueryBuilder, context);
+            } else {
+                testDoToQuery(queryBuilder, context);
             }
         });
     }
@@ -301,8 +246,8 @@ private void testDoToQuery(SparseVectorQueryBuilder queryBuilder, SearchExecutio
         assertTrue(query instanceof SparseVectorQueryWrapper);
         var sparseQuery = (SparseVectorQueryWrapper) query;
 
-        // check if we have explicit pruning, or pruning via the index_options
-        if (queryBuilder.shouldPruneTokens() || currentTestHasIndexOptions()) {
+        // check if we have explicit pruning or implicit pruning=true
+        if (queryBuilder.shouldPruneTokens() || indexVersionSupportsIndexOptions()) {
             // It's possible that all documents were pruned for aggressive pruning configurations
             assertTrue(sparseQuery.getTermsQuery() instanceof BooleanQuery || sparseQuery.getTermsQuery() instanceof MatchNoDocsQuery);
         } else {
@@ -401,48 +346,18 @@ public void testThatWeCorrectlyRewriteQueryIntoVectors() {
         assertNotNull(((SparseVectorQueryBuilder) rewrittenQueryBuilder).getQueryVectors());
     }
 
-    @InjectSparseVectorIndexOptions
-    public void testItUsesIndexOptionsDefaults() throws IOException {
-        withSearchIndex((context) -> {
-            try {
-                SparseVectorQueryBuilder builder = createTestQueryBuilder(null);
-                assertFalse(builder.shouldPruneTokens());
-                testDoToQuery(builder, context);
-            } catch (IOException ex) {
-                throw new RuntimeException(ex);
-            }
-        });
-    }
+    private boolean indexVersionSupportsIndexOptions() {
+        if (indexVersionToTest.onOrAfter(IndexVersions.SPARSE_VECTOR_PRUNING_INDEX_OPTIONS_SUPPORT)) {
+            return true;
+        }
 
-    @InjectSparseVectorIndexOptions
-    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);
-            }
-        });
-    }
+        if (indexVersionToTest.between(
+            IndexVersions.SPARSE_VECTOR_PRUNING_INDEX_OPTIONS_SUPPORT_BACKPORT_8_X,
+            IndexVersions.UPGRADE_TO_LUCENE_10_0_0
+        )) {
+            return true;
+        }
 
-    @InjectSparseVectorIndexOptions
-    public void testToQueryRewriteWithIndexOptions() throws IOException {
-        withSearchIndex((context) -> {
-            SparseVectorQueryBuilder queryBuilder = createTestQueryBuilder(null);
-            try {
-                if (queryBuilder.getQueryVectors() == null) {
-                    QueryBuilder rewrittenQueryBuilder = rewriteAndFetch(queryBuilder, context);
-                    assertTrue(rewrittenQueryBuilder instanceof SparseVectorQueryBuilder);
-                    testDoToQuery((SparseVectorQueryBuilder) rewrittenQueryBuilder, context);
-                } else {
-                    testDoToQuery(queryBuilder, context);
-                }
-            } catch (IOException e) {
-                throw new RuntimeException(e);
-            }
-        });
+        return false;
     }
 }

x-pack/plugin/ml/qa/multi-cluster-tests-with-security/src/test/resources/rest-api-spec/test/multi_cluster/50_sparse_vector.yml

@@ -273,7 +273,6 @@ teardown:
         index: sparse_vector_pruning_test
 
   - not_exists: sparse_vector_pruning_test.mappings.properties.ml.properties.tokens.index_options
-  - not_exists: sparse_vector_pruning_test.mappings.properties.embeddings.index_options
 
 ---
 "Check sparse_vector token pruning index_options prune missing do not allow config":
@@ -554,3 +553,60 @@ teardown:
   - match: { hits.hits.0._id: "1" }
   - match: { hits.hits.1._id: "3" }
   - match: { hits.hits.2._id: "2" }
+
+---
+"Check sparse_vector should prune by default":
+
+  - requires:
+      cluster_features: 'sparse_vector.index_options_supported'
+      reason: "sparse_vector token pruning index options added support in 8.19"
+  - skip:
+      features: headers
+
+  - do:
+      headers:
+        Authorization: "Basic dGVzdF91c2VyOngtcGFjay10ZXN0LXBhc3N3b3Jk" #test_user credentials
+        Content-Type: application/json
+      indices.create:
+        index: test-sparse-vector-pruning-default
+        body:
+          mappings:
+            properties:
+              content_embedding:
+                type: sparse_vector
+
+  - match: { acknowledged: true }
+
+  - do:
+      headers:
+        Authorization: "Basic dGVzdF91c2VyOngtcGFjay10ZXN0LXBhc3N3b3Jk" #test_user credentials
+        Content-Type: application/json
+      bulk:
+        index: test-sparse-vector-pruning-default
+        refresh: true
+        body: |
+          {"index": { "_id": "1" }}
+          {"content_embedding":{"cheese": 2.671405,"is": 0.11809908,"comet": 0.26088917}}
+          {"index": { "_id": "2" }}
+          {"content_embedding":{"planet": 2.3438394,"is": 0.54600334,"astronomy": 0.36015007,"moon": 0.20022368}}
+          {"index": { "_id": "3" }}
+          {"content_embedding":{"is": 0.6891394,"globe": 0.484035,"ocean": 0.080102935,"underground": 0.053516876}}
+
+  - do:
+      search:
+        index: test-sparse-vector-pruning-default
+        body:
+          query:
+            sparse_vector:
+              field: content_embedding
+              query_vector:
+                cheese: 0.5
+                comet: 0.5
+                globe: 0.484035
+                ocean: 0.080102935
+                underground: 0.053516876
+                is: 0.54600334
+
+  - match: { hits.total.value: 2 }
+  - match: { hits.hits.0._id: "1" }
+  - match: { hits.hits.1._id: "3" }