GPU: Test feature disabling and document mutations (#138966)

- Add test for GPU-to-non-GPU transition (indexing on GPU,
  with continuous search and indexing on CPU)
- Add test for deletes/updates
- Support configurable similarity metrics in test setup

Author: mayya-sharipova

x-pack/plugin/gpu/src/internalClusterTest/java/org/elasticsearch/xpack/gpu/GPUIndexIT.java

@@ -18,6 +18,7 @@
 import org.elasticsearch.search.vectors.KnnSearchBuilder;
 import org.elasticsearch.search.vectors.VectorData;
 import org.elasticsearch.test.ESIntegTestCase;
+import org.junit.After;
 import org.junit.Assert;
 import org.junit.BeforeClass;
 
@@ -35,17 +36,27 @@
 @LuceneTestCase.SuppressCodecs("*") // use our custom codec
 public class GPUIndexIT extends ESIntegTestCase {
 
+    private static boolean isGpuIndexingFeatureAllowed = true;
+
+    private String similarity;
+
     public static class TestGPUPlugin extends GPUPlugin {
         public TestGPUPlugin() {
             super(Settings.builder().put("vectors.indexing.use_gpu", GpuMode.TRUE.name()).build());
         }
 
         @Override
         protected boolean isGpuIndexingFeatureAllowed() {
-            return true;
+            return GPUIndexIT.isGpuIndexingFeatureAllowed;
         }
     }
 
+    @After
+    public void reset() {
+        isGpuIndexingFeatureAllowed = true;
+        similarity = null;
+    }
+
     @Override
     protected Collection<Class<? extends Plugin>> nodePlugins() {
         return List.of(TestGPUPlugin.class);
@@ -70,6 +81,36 @@ public void testBasic() {
         assertSearch(indexName, randomFloatVector(dims), totalDocs);
     }
 
+    public void testSearchAndIndexAfterDisablingGpu() {
+        String indexName = "index1";
+        final int dims = randomIntBetween(4, 128);
+        final int numDocs = randomIntBetween(1, 500);
+        createIndex(indexName, dims, false);
+        ensureGreen();
+
+        indexDocs(indexName, numDocs, dims, 0);
+        refresh();
+
+        // Disable GPU usage via feature flag (simulating missing license)
+        isGpuIndexingFeatureAllowed = false;
+        ensureGreen();
+
+        assertAcked(indicesAdmin().prepareClose(indexName).get());
+        assertAcked(indicesAdmin().prepareOpen(indexName).get());
+        ensureGreen();
+
+        assertSearch(indexName, randomFloatVector(dims), numDocs);
+
+        // Add more data to the index
+        final int additionalDocs = randomIntBetween(1, 100);
+        indexDocs(indexName, additionalDocs, dims, numDocs);
+        refresh();
+        final int totalDocs = numDocs + additionalDocs;
+
+        // Perform another search with the additional data
+        assertSearch(indexName, randomFloatVector(dims), totalDocs);
+    }
+
     public void testSortedIndexReturnsSameResultsAsUnsorted() {
         String indexName1 = "index_unsorted";
         String indexName2 = "index_sorted";
@@ -193,6 +234,58 @@ public void testSortedIndexReturnsSameResultsAsUnsorted() {
         }
     }
 
+    public void testDeletesUpdates() {
+        String indexName = "index_deletes_updates";
+        final int dims = randomIntBetween(4, 128);
+        createIndex(indexName, dims, false);
+
+        final int numDocs = randomIntBetween(700, 1000);
+        indexDocs(indexName, numDocs, dims, 0);
+        refresh();
+
+        // Perform random updates and deletes
+        final int numOperations = randomIntBetween(10, 50);
+        BulkRequestBuilder bulkRequest = client().prepareBulk();
+        for (int i = 0; i < numOperations; i++) {
+            int docId = randomIntBetween(0, numDocs - 1);
+            if (randomBoolean()) {
+                bulkRequest.add(
+                    prepareIndex(indexName).setId(String.valueOf(docId))
+                        .setSource("my_vector", randomFloatVector(dims), "my_keyword", String.valueOf(randomIntBetween(1, numDocs)))
+                );
+            } else {
+                bulkRequest.add(client().prepareDelete(indexName, String.valueOf(docId)));
+            }
+        }
+        BulkResponse bulkResponse = bulkRequest.get();
+        assertFalse("Bulk request failed: " + bulkResponse.buildFailureMessage(), bulkResponse.hasFailures());
+        refresh();
+
+        // Assert that approximate and exact searches return same sets of results
+        float[] queryVector = randomFloatVector(dims);
+        int k = 10;
+        int numCandidates = k * 10;
+
+        var approxSearchResponse = prepareSearch(indexName).setSize(k)
+            .setFetchSource(false)
+            .setKnnSearch(List.of(new KnnSearchBuilder("my_vector", queryVector, k, numCandidates, null, null, null)))
+            .get();
+
+        var exactSearchResponse = prepareSearch(indexName).setSize(k)
+            .setFetchSource(false)
+            .setQuery(new ExactKnnQueryBuilder(VectorData.fromFloats(queryVector), "my_vector", null))
+            .get();
+
+        try {
+            SearchHit[] approxHits = approxSearchResponse.getHits().getHits();
+            SearchHit[] exactHits = exactSearchResponse.getHits().getHits();
+            assertAtLeastNOutOfKMatches(approxHits, exactHits, k - 3, k);
+        } finally {
+            approxSearchResponse.decRef();
+            exactSearchResponse.decRef();
+        }
+    }
+
     public void testInt8HnswMaxInnerProductProductFails() {
         String indexName = "index_int8_max_inner_product_fails";
         final int dims = randomIntBetween(4, 128);
@@ -222,7 +315,7 @@ public void testInt8HnswMaxInnerProductProductFails() {
         // Attempt to index a document and expect it to fail
         IllegalArgumentException ex = expectThrows(
             IllegalArgumentException.class,
-            () -> client().prepareIndex(indexName).setId("1").setSource("my_vector", randomFloatVector(dims)).get()
+            () -> client().prepareIndex(indexName).setId("1").setSource("my_vector", randomNonUnitFloatVector(dims)).get()
         );
         assertThat(
             ex.getMessage(),
@@ -237,14 +330,18 @@ private void createIndex(String indexName, int dims, boolean sorted) {
             settings.put("index.sort.field", "my_keyword");
         }
 
+        if (similarity == null) {
+            similarity = randomFrom("dot_product", "l2_norm", "cosine");
+        }
+
         String type = randomFrom("hnsw", "int8_hnsw");
         String mapping = String.format(Locale.ROOT, """
             {
               "properties": {
                 "my_vector": {
                   "type": "dense_vector",
                   "dims": %d,
-                  "similarity": "l2_norm",
+                  "similarity": "%s",
                   "index_options": {
                     "type": "%s"
                   }
@@ -254,7 +351,7 @@ private void createIndex(String indexName, int dims, boolean sorted) {
                 }
               }
             }
-            """, dims, type);
+            """, dims, similarity, type);
         assertAcked(prepareCreate(indexName).setSettings(settings.build()).setMapping(mapping));
         ensureGreen();
     }
@@ -264,8 +361,8 @@ private void indexDocs(String indexName, int numDocs, int dims, int startDoc) {
         for (int i = 0; i < numDocs; i++) {
             String id = String.valueOf(startDoc + i);
             String keywordValue = String.valueOf(numDocs - i);
-            var indexRequest = prepareIndex(indexName).setId(id)
-                .setSource("my_vector", randomFloatVector(dims), "my_keyword", keywordValue);
+            float[] vector = randomFloatVector(dims);
+            var indexRequest = prepareIndex(indexName).setId(id).setSource("my_vector", vector, "my_keyword", keywordValue);
             bulkRequest.add(indexRequest);
         }
         BulkResponse bulkResponse = bulkRequest.get();
@@ -284,14 +381,35 @@ private void assertSearch(String indexName, float[] queryVector, int totalDocs)
         );
     }
 
-    private static float[] randomFloatVector(int dims) {
+    private static float[] randomNonUnitFloatVector(int dims) {
         float[] vector = new float[dims];
         for (int i = 0; i < dims; i++) {
             vector[i] = randomFloat();
         }
         return vector;
     }
 
+    private static float[] randomUnitVector(int dims) {
+        float[] vector = new float[dims];
+        double sumSquares = 0.0;
+        for (int i = 0; i < dims; i++) {
+            vector[i] = randomFloat() * 2 - 1; // Generate values between -1 and 1 for better distribution
+            sumSquares += vector[i] * vector[i];
+        }
+        float magnitude = (float) Math.sqrt(sumSquares);
+        if (magnitude > 0) {
+            for (int i = 0; i < dims; i++) {
+                vector[i] /= magnitude;
+            }
+        }
+        return vector;
+    }
+
+    private float[] randomFloatVector(int dims) {
+        boolean useUnitVectors = "dot_product".equals(similarity);
+        return useUnitVectors ? randomUnitVector(dims) : randomNonUnitFloatVector(dims);
+    }
+
     /**
      * Asserts that at least N out of K hits have matching IDs between two result sets.
      */