Implement working with sorted index

Author: mayya-sharipova

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

@@ -12,10 +12,12 @@
 import org.elasticsearch.action.bulk.BulkResponse;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.plugins.Plugin;
+import org.elasticsearch.search.SearchHit;
 import org.elasticsearch.search.vectors.KnnSearchBuilder;
 import org.elasticsearch.test.ESIntegTestCase;
 import org.elasticsearch.xpack.gpu.GPUPlugin;
 import org.elasticsearch.xpack.gpu.GPUSupport;
+import org.junit.Assert;
 
 import java.util.Collection;
 import java.util.List;
@@ -34,40 +36,102 @@ protected Collection<Class<? extends Plugin>> nodePlugins() {
 
     public void testBasic() {
         assumeTrue("cuvs not supported", GPUSupport.isSupported(false));
+        String indexName = "index1";
         final int dims = randomIntBetween(4, 128);
         final int[] numDocs = new int[] { randomIntBetween(1, 100), 1, 2, randomIntBetween(1, 100) };
-        createIndex(dims);
+        createIndex(indexName, dims, false);
         int totalDocs = 0;
         for (int i = 0; i < numDocs.length; i++) {
-            indexDocs(numDocs[i], dims, i * 100);
+            indexDocs(indexName, numDocs[i], dims, i * 100);
             totalDocs += numDocs[i];
         }
         refresh();
-        assertSearch(randomFloatVector(dims), totalDocs);
+        assertSearch(indexName, randomFloatVector(dims), totalDocs);
+    }
+
+    public void testSortedIndexReturnsSameResultsAsUnsorted() {
+        assumeTrue("cuvs not supported", GPUSupport.isSupported(false));
+        String indexName1 = "index_unsorted";
+        String indexName2 = "index_sorted";
+        final int dims = randomIntBetween(4, 128);
+        createIndex(indexName1, dims, false);
+        createIndex(indexName2, dims, true);
+
+        final int[] numDocs = new int[] { randomIntBetween(50, 100), randomIntBetween(50, 100) };
+        for (int i = 0; i < numDocs.length; i++) {
+            BulkRequestBuilder bulkRequest1 = client().prepareBulk();
+            BulkRequestBuilder bulkRequest2 = client().prepareBulk();
+            for (int j = 0; j < numDocs[i]; j++) {
+                String id = String.valueOf(i * 100 + j);
+                String keywordValue = String.valueOf(numDocs[i] - j);
+                float[] vector = randomFloatVector(dims);
+                bulkRequest1.add(prepareIndex(indexName1).setId(id).setSource("my_vector", vector, "my_keyword", keywordValue));
+                bulkRequest2.add(prepareIndex(indexName2).setId(id).setSource("my_vector", vector, "my_keyword", keywordValue));
+            }
+            BulkResponse bulkResponse1 = bulkRequest1.get();
+            assertFalse("Bulk request failed: " + bulkResponse1.buildFailureMessage(), bulkResponse1.hasFailures());
+            BulkResponse bulkResponse2 = bulkRequest2.get();
+            assertFalse("Bulk request failed: " + bulkResponse2.buildFailureMessage(), bulkResponse2.hasFailures());
+        }
+        refresh();
+
+        float[] queryVector = randomFloatVector(dims);
+        int k = 10;
+        int numCandidates = k * 10;
+
+        var searchResponse1 = prepareSearch(indexName1).setSize(k)
+            .setFetchSource(false)
+            .addFetchField("my_keyword")
+            .setKnnSearch(List.of(new KnnSearchBuilder("my_vector", queryVector, k, numCandidates, null, null)))
+            .get();
+
+        var searchResponse2 = prepareSearch(indexName2).setSize(k)
+            .setFetchSource(false)
+            .addFetchField("my_keyword")
+            .setKnnSearch(List.of(new KnnSearchBuilder("my_vector", queryVector, k, numCandidates, null, null)))
+            .get();
+
+        try {
+            SearchHit[] hits1 = searchResponse1.getHits().getHits();
+            SearchHit[] hits2 = searchResponse2.getHits().getHits();
+            Assert.assertEquals(hits1.length, hits2.length);
+            for (int i = 0; i < hits1.length; i++) {
+                Assert.assertEquals(hits1[i].getId(), hits2[i].getId());
+                Assert.assertEquals((String) hits1[i].field("my_keyword").getValue(), (String) hits2[i].field("my_keyword").getValue());
+                Assert.assertEquals(hits1[i].getScore(), hits2[i].getScore(), 0.0001f);
+            }
+        } finally {
+            searchResponse1.decRef();
+            searchResponse2.decRef();
+        }
     }
 
     public void testSearchWithoutGPU() {
         assumeTrue("cuvs not supported", GPUSupport.isSupported(false));
+        String indexName = "index1";
         final int dims = randomIntBetween(4, 128);
         final int numDocs = randomIntBetween(1, 500);
-        createIndex(dims);
+        createIndex(indexName, dims, false);
         ensureGreen();
 
-        indexDocs(numDocs, dims, 0);
+        indexDocs(indexName, numDocs, dims, 0);
         refresh();
 
         // update settings to disable GPU usage
         Settings.Builder settingsBuilder = Settings.builder().put("index.vectors.indexing.use_gpu", false);
-        assertAcked(client().admin().indices().prepareUpdateSettings("foo-index").setSettings(settingsBuilder.build()));
+        assertAcked(client().admin().indices().prepareUpdateSettings(indexName).setSettings(settingsBuilder.build()));
         ensureGreen();
-        assertSearch(randomFloatVector(dims), numDocs);
+        assertSearch(indexName, randomFloatVector(dims), numDocs);
     }
 
-    private void createIndex(int dims) {
+    private void createIndex(String indexName, int dims, boolean sorted) {
         var settings = Settings.builder().put(indexSettings());
         settings.put("index.number_of_shards", 1);
         settings.put("index.vectors.indexing.use_gpu", true);
-        assertAcked(prepareCreate("foo-index").setSettings(settings.build()).setMapping(String.format(Locale.ROOT, """
+        if (sorted) {
+            settings.put("index.sort.field", "my_keyword");
+        }
+        assertAcked(prepareCreate(indexName).setSettings(settings.build()).setMapping(String.format(Locale.ROOT, """
                 {
                   "properties": {
                     "my_vector": {
@@ -77,28 +141,36 @@ private void createIndex(int dims) {
                       "index_options": {
                         "type": "hnsw"
                       }
+                    },
+                    "my_keyword": {
+                      "type": "keyword"
                     }
                   }
                 }
             """, dims)));
         ensureGreen();
     }
 
-    private void indexDocs(int numDocs, int dims, int startDoc) {
+    private void indexDocs(String indexName, int numDocs, int dims, int startDoc) {
         BulkRequestBuilder bulkRequest = client().prepareBulk();
         for (int i = 0; i < numDocs; i++) {
             String id = String.valueOf(startDoc + i);
-            bulkRequest.add(prepareIndex("foo-index").setId(id).setSource("my_vector", randomFloatVector(dims)));
+            String keywordValue = String.valueOf(numDocs - i);
+            var indexRequest = prepareIndex(indexName).setId(id)
+                .setSource("my_vector", randomFloatVector(dims), "my_keyword", keywordValue);
+            bulkRequest.add(indexRequest);
         }
         BulkResponse bulkResponse = bulkRequest.get();
         assertFalse("Bulk request failed: " + bulkResponse.buildFailureMessage(), bulkResponse.hasFailures());
     }
 
-    private void assertSearch(float[] queryVector, int totalDocs) {
+    private void assertSearch(String indexName, float[] queryVector, int totalDocs) {
         int k = Math.min(randomIntBetween(1, 20), totalDocs);
         int numCandidates = k * 10;
         assertNoFailuresAndResponse(
-            prepareSearch("foo-index").setSize(k)
+            prepareSearch(indexName).setSize(k)
+                .setFetchSource(false)
+                .addFetchField("my_keyword")
                 .setKnnSearch(List.of(new KnnSearchBuilder("my_vector", queryVector, k, numCandidates, null, null))),
             response -> {
                 assertEquals("Expected k hits to be returned", k, response.getHits().getHits().length);

x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/GPUToHNSWVectorsWriter.java

@@ -221,9 +221,10 @@ private void writeField(FieldWriter fieldWriter) throws IOException {
     }
 
     private void writeSortingField(FieldWriter fieldData, Sorter.DocMap sortMap) throws IOException {
-        // TODO: implement writing sorted field when we can access cagra index through MemorySegment
-        // as we need random access to neighbors in the graph.
-        throw new UnsupportedOperationException("Writing field with index sorted needs to be implemented.");
+        // The flatFieldVectorsWriter's flush method, called before this, has already sorted the vectors according to the sortMap.
+        // We can now treat them as a simple, sorted list of vectors.
+        float[][] vectors = fieldData.flatFieldVectorsWriter.getVectors().toArray(float[][]::new);
+        writeFieldInternal(fieldData.fieldInfo, DatasetOrVectors.fromArray(vectors));
     }
 
     private void writeFieldInternal(FieldInfo fieldInfo, DatasetOrVectors datasetOrVectors) throws IOException {