Add new kNN search endpoint

rest-api-spec/src/main/resources/rest-api-spec/api/knn_search.json

@@ -0,0 +1,47 @@
+{
+  "knn_search":{
+    "documentation":{
+      "url":"https://www.elastic.co/guide/en/elasticsearch/reference/master/search-search.html",
+      "description":"Performs a kNN search."
+    },
+    "stability":"experimental",
+    "visibility":"public",
+    "headers":{
+      "accept": [ "application/json"],
+      "content_type": ["application/json"]
+    },
+    "url":{
+      "paths":[
+        {
+          "path":"/_knn_search",
+          "methods":[
+            "GET",
+            "POST"
+          ]
+        },
+        {
+          "path":"/{index}/_knn_search",
+          "methods":[
+            "GET",
+            "POST"
+          ],
+          "parts":{
+            "index":{
+              "type":"list",
+              "description":"A comma-separated list of index names to search; use `_all` or empty string to perform the operation on all indices"
+            }
+          }
+        }
+      ]
+    },
+    "params": {
+      "routing":{
+        "type":"list",
+        "description":"A comma-separated list of specific routing values"
+      }
+    },
+    "body":{
+      "description":"The search definition"
+    }
+  }
+}

server/src/main/java/org/elasticsearch/search/builder/SearchSourceBuilder.java

@@ -16,13 +16,7 @@
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.common.logging.DeprecationLogger;
-import org.elasticsearch.xcontent.ParseField;
-import org.elasticsearch.xcontent.ToXContentFragment;
-import org.elasticsearch.xcontent.ToXContentObject;
-import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentHelper;
-import org.elasticsearch.xcontent.XContentParser;
-import org.elasticsearch.xcontent.XContentType;
 import org.elasticsearch.core.Booleans;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.RestApiVersion;
@@ -49,6 +43,12 @@
 import org.elasticsearch.search.sort.SortBuilders;
 import org.elasticsearch.search.sort.SortOrder;
 import org.elasticsearch.search.suggest.SuggestBuilder;
+import org.elasticsearch.xcontent.ParseField;
+import org.elasticsearch.xcontent.ToXContentFragment;
+import org.elasticsearch.xcontent.ToXContentObject;
+import org.elasticsearch.xcontent.XContentBuilder;
+import org.elasticsearch.xcontent.XContentParser;
+import org.elasticsearch.xcontent.XContentType;
 
 import java.io.IOException;
 import java.util.ArrayList;

test/framework/src/main/java/org/elasticsearch/search/RandomSearchRequestGenerator.java

@@ -14,6 +14,7 @@
 import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.text.Text;
 import org.elasticsearch.core.TimeValue;
+import org.elasticsearch.search.fetch.subphase.FieldAndFormat;
 import org.elasticsearch.xcontent.DeprecationHandler;
 import org.elasticsearch.xcontent.NamedXContentRegistry;
 import org.elasticsearch.xcontent.XContentBuilder;
@@ -174,7 +175,18 @@ public static SearchSourceBuilder randomSearchSourceBuilder(
         if (randomBoolean()) {
             int numFields = randomInt(5);
             for (int i = 0; i < numFields; i++) {
-                builder.fetchField(randomAlphaOfLengthBetween(5, 10));
+                String field = randomAlphaOfLengthBetween(5, 10);
+                String format = randomBoolean() ? randomAlphaOfLengthBetween(5, 10) : null;
+                builder.fetchField(new FieldAndFormat(field, format));
+            }
+        }
+
+        if (randomBoolean()) {
+            int numFields = randomInt(5);
+            for (int i = 0; i < numFields; i++) {
+                String field = randomAlphaOfLengthBetween(5, 10);
+                String format = randomBoolean() ? randomAlphaOfLengthBetween(5, 10) : null;
+                builder.docValueField(field, format);
             }
         }
 

x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/vectors/40_knn_search.yml

@@ -0,0 +1,90 @@
+setup:
+  - do:
+      indices.create:
+        index: test
+        body:
+          settings:
+            number_of_replicas: 0
+          mappings:
+            properties:
+              name:
+                type: keyword
+              vector:
+                type: dense_vector
+                dims: 5
+                index: true
+                similarity: l2_norm
+  - do:
+      index:
+        index: test
+        body:
+          name: cow.jpg
+          vector: [230.0, 300.33, -34.8988, 15.555, -200.0]
+
+  - do:
+      index:
+        index: test
+        id: 2
+        body:
+          name: moose.jpg
+          vector: [-0.5, 100.0, -13, 14.8, -156.0]
+
+  - do:
+      index:
+        index: test
+        id: 3
+        body:
+          name: rabbit.jpg
+          vector: [0.5, 111.3, -13.0, 14.8, -156.0]
+
+  - do:
+      indices.refresh: {}
+
+---
+"Basic kNN search":
+  - do:
+      knn_search:
+        index: test
+        body:
+          fields: [ "name" ]
+          knn:
+            field: vector
+            query_vector: [-0.5, 90.0, -10, 14.8, -156.0]
+            k: 2
+            num_candidates: 3
+
+  - match: {hits.hits.0._id: "2"}
+  - match: {hits.hits.0.fields.name.0: "moose.jpg"}
+
+  - match: {hits.hits.1._id: "3"}
+  - match: {hits.hits.1.fields.name.0: "rabbit.jpg"}
+
+---
+"Basic kNN search with no index":
+  - do:
+      knn_search:
+        body:
+          fields: [ "name" ]
+          knn:
+            field: vector
+            query_vector: [-0.5, 90.0, -10, 14.8, -156.0]
+            k: 1
+            num_candidates: 1
+
+  - match: {hits.hits.0._id: "2"}
+  - match: {hits.hits.0.fields.name.0: "moose.jpg"}
+
+---
+"Direct knn queries are disallowed":
+  - do:
+      catch: bad_request
+      search:
+        rest_total_hits_as_int: true
+        index: test-index
+        body:
+          query:
+            knn:
+              field: vector
+              query_vector: [ -0.5, 90.0, -10, 14.8, -156.0 ]
+              num_candidates: 1
+  - match: { error.root_cause.0.type: "illegal_argument_exception" }

x-pack/plugin/vectors/src/main/java/org/elasticsearch/xpack/vectors/DenseVectorPlugin.java

@@ -7,17 +7,31 @@
 
 package org.elasticsearch.xpack.vectors;
 
+import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
+import org.elasticsearch.cluster.node.DiscoveryNodes;
+import org.elasticsearch.common.settings.ClusterSettings;
+import org.elasticsearch.common.settings.IndexScopedSettings;
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.common.settings.SettingsFilter;
 import org.elasticsearch.index.mapper.Mapper;
+import org.elasticsearch.plugins.ActionPlugin;
 import org.elasticsearch.plugins.MapperPlugin;
 import org.elasticsearch.plugins.Plugin;
+import org.elasticsearch.plugins.SearchPlugin;
+import org.elasticsearch.rest.RestController;
+import org.elasticsearch.rest.RestHandler;
+import org.elasticsearch.xpack.vectors.action.RestKnnSearchAction;
 import org.elasticsearch.xpack.vectors.mapper.DenseVectorFieldMapper;
 import org.elasticsearch.xpack.vectors.mapper.SparseVectorFieldMapper;
+import org.elasticsearch.xpack.vectors.query.KnnVectorQueryBuilder;
 
 import java.util.Collections;
 import java.util.LinkedHashMap;
+import java.util.List;
 import java.util.Map;
+import java.util.function.Supplier;
 
-public class DenseVectorPlugin extends Plugin implements MapperPlugin {
+public class DenseVectorPlugin extends Plugin implements ActionPlugin, MapperPlugin, SearchPlugin {
 
     public DenseVectorPlugin() { }
 
@@ -28,4 +42,26 @@ public Map<String, Mapper.TypeParser> getMappers() {
         mappers.put(SparseVectorFieldMapper.CONTENT_TYPE, SparseVectorFieldMapper.PARSER);
         return Collections.unmodifiableMap(mappers);
     }
+
+    @Override
+    public List<RestHandler> getRestHandlers(
+        Settings settings,
+        RestController restController,
+        ClusterSettings clusterSettings,
+        IndexScopedSettings indexScopedSettings,
+        SettingsFilter settingsFilter,
+        IndexNameExpressionResolver indexNameExpressionResolver,
+        Supplier<DiscoveryNodes> nodesInCluster
+    ) {
+        return List.of(new RestKnnSearchAction());
+    }
+
+    @Override
+    public List<QuerySpec<?>> getQueries() {
+        // This query is only meant to be used internally, and not passed to the _search endpoint
+        return List.of(new QuerySpec<>(KnnVectorQueryBuilder.NAME, KnnVectorQueryBuilder::new,
+            parser -> {
+                throw new IllegalArgumentException("[knn] queries cannot be provided directly, use the [_knn_search] endpoint instead");
+            }));
+    }
 }