feat: Add kNN Query (#198)

foxstarius · KennyLindahl · Andreas Franzon · web-flow · commit 78de1797fd19 · 2024-05-05T17:38:57.000+05:30
kNN search was added to Elasticsearch in v8.0

Co-authored-by: kennylindahl &lt;haxxblaster@gmail.com&gt;
Co-authored-by: Andreas Franzon &lt;andreas.franzon@schibsted.com&gt;
diff --git a/src/core/index.js b/src/core/index.js
@@ -8,6 +8,8 @@ exports.Aggregation = require('./aggregation');
 
 exports.Query = require('./query');
 
+exports.KNN = require('./knn');
+
 exports.Suggester = require('./suggester');
 
 exports.Script = require('./script');
diff --git a/src/core/knn.js b/src/core/knn.js
@@ -0,0 +1,138 @@
+'use strict';
+
+const { recursiveToJSON, checkType } = require('./util');
+const Query = require('./query');
+
+/**
+ * Class representing a k-Nearest Neighbors (k-NN) query.
+ * This class extends the Query class to support the specifics of k-NN search, including setting up the field,
+ * query vector, number of neighbors (k), and number of candidates.
+ *
+ * @example
+ * const qry = esb.kNN('my_field', 100, 1000).vector([1,2,3]);
+ * const qry = esb.kNN('my_field', 100, 1000).queryVectorBuilder('model_123', 'Sample model text');
+ *
+ * NOTE: kNN search was added to Elasticsearch in v8.0
+ *
+ * [Elasticsearch reference](https://www.elastic.co/guide/en/elasticsearch/reference/current/knn-search.html)
+ */
+class KNN {
+    // eslint-disable-next-line require-jsdoc
+    constructor(field, k, numCandidates) {
+        if (k > numCandidates)
+            throw new Error('KNN numCandidates cannot be less than k');
+        this._body = {};
+        this._body.field = field;
+        this._body.k = k;
+        this._body.filter = [];
+        this._body.num_candidates = numCandidates;
+    }
+
+    /**
+     * Sets the query vector for the k-NN search.
+     * @param {Array<number>} vector - The query vector.
+     * @returns {KNN} Returns the instance of KNN for method chaining.
+     */
+    queryVector(vector) {
+        if (this._body.query_vector_builder)
+            throw new Error(
+                'cannot provide both query_vector_builder and query_vector'
+            );
+        this._body.query_vector = vector;
+        return this;
+    }
+
+    /**
+     * Sets the query vector builder for the k-NN search.
+     * This method configures a query vector builder using a specified model ID and model text.
+     * It's important to note that either a direct query vector or a query vector builder can be
+     * provided, but not both.
+     *
+     * @param {string} modelId - The ID of the model to be used for generating the query vector.
+     * @param {string} modelText - The text input based on which the query vector is generated.
+     * @returns {KNN} Returns the instance of KNN for method chaining.
+     * @throws {Error} Throws an error if both query_vector_builder and query_vector are provided.
+     *
+     * @example
+     * let knn = new esb.KNN().queryVectorBuilder('model_123', 'Sample model text');
+     */
+    queryVectorBuilder(modelId, modelText) {
+        if (this._body.query_vector)
+            throw new Error(
+                'cannot provide both query_vector_builder and query_vector'
+            );
+        this._body.query_vector_builder = {
+            text_embeddings: {
+                model_id: modelId,
+                model_text: modelText
+            }
+        };
+        return this;
+    }
+
+    /**
+     * Adds one or more filter queries to the k-NN search.
+     *
+     * This method is designed to apply filters to the k-NN search. It accepts either a single
+     * query or an array of queries. Each query acts as a filter, refining the search results
+     * according to the specified conditions. These queries must be instances of the `Query` class.
+     * If any provided query is not an instance of `Query`, a TypeError is thrown.
+     *
+     * @param {Query|Query[]} queries - A single `Query` instance or an array of `Query` instances for filtering.
+     * @returns {KNN} Returns `this` to allow method chaining.
+     * @throws {TypeError} If any of the provided queries is not an instance of `Query`.
+     *
+     * @example
+     * let knn = new esb.KNN().filter(new esb.TermQuery('field', 'value')); // Applying a single filter query
+     *
+     * @example
+     * let knn = new esb.KNN().filter([
+     *     new esb.TermQuery('field1', 'value1'),
+     *     new esb.TermQuery('field2', 'value2')
+     * ]); // Applying multiple filter queries
+     */
+    filter(queries) {
+        const queryArray = Array.isArray(queries) ? queries : [queries];
+        queryArray.forEach(query => {
+            checkType(query, Query);
+            this._body.filter.push(query);
+        });
+        return this;
+    }
+
+    /**
+     * Sets the field to perform the k-NN search on.
+     * @param {number} boost - The number of the boost
+     * @returns {KNN} Returns the instance of KNN for method chaining.
+     */
+    boost(boost) {
+        this._body.boost = boost;
+        return this;
+    }
+
+    /**
+     * Sets the field to perform the k-NN search on.
+     * @param {number} similarity - The number of the similarity
+     * @returns {KNN} Returns the instance of KNN for method chaining.
+     */
+    similarity(similarity) {
+        this._body.similarity = similarity;
+        return this;
+    }
+
+    /**
+     * Override default `toJSON` to return DSL representation for the `query`
+     *
+     * @override
+     * @returns {Object} returns an Object which maps to the elasticsearch query DSL
+     */
+    toJSON() {
+        if (!this._body.query_vector && !this._body.query_vector_builder)
+            throw new Error(
+                'either query_vector_builder or query_vector must be provided'
+            );
+        return recursiveToJSON(this._body);
+    }
+}
+
+module.exports = KNN;
diff --git a/src/core/request-body-search.js b/src/core/request-body-search.js
@@ -10,7 +10,8 @@ const Query = require('./query'),
     Rescore = require('./rescore'),
     Sort = require('./sort'),
     Highlight = require('./highlight'),
-    InnerHits = require('./inner-hits');
+    InnerHits = require('./inner-hits'),
+    KNN = require('./knn');
 
 const { checkType, setDefault, recursiveToJSON } = require('./util');
 const RuntimeField = require('./runtime-field');
@@ -70,6 +71,7 @@ class RequestBodySearch {
     constructor() {
         // Maybe accept some optional parameter?
         this._body = {};
+        this._knn = [];
         this._aggs = [];
         this._suggests = [];
         this._suggestText = null;
@@ -88,6 +90,21 @@ class RequestBodySearch {
         return this;
     }
 
+    /**
+     * Sets knn on the search request body.
+     *
+     * @param {Knn|Knn[]} knn
+     * @returns {RequestBodySearch} returns `this` so that calls can be chained.
+     */
+    kNN(knn) {
+        const knns = Array.isArray(knn) ? knn : [knn];
+        knns.forEach(_knn => {
+            checkType(_knn, KNN);
+            this._knn.push(_knn);
+        });
+        return this;
+    }
+
     /**
      * Sets aggregation on the request body.
      * Alias for method `aggregation`
@@ -867,6 +884,12 @@ class RequestBodySearch {
     toJSON() {
         const dsl = recursiveToJSON(this._body);
 
+        if (!isEmpty(this._knn))
+            dsl.knn =
+                this._knn.length == 1
+                    ? recMerge(this._knn)
+                    : this._knn.map(knn => recursiveToJSON(knn));
+
         if (!isEmpty(this._aggs)) dsl.aggs = recMerge(this._aggs);
 
         if (!isEmpty(this._suggests) || !isNil(this._suggestText)) {
diff --git a/src/index.d.ts b/src/index.d.ts
@@ -18,6 +18,13 @@ declare namespace esb {
          */
         query(query: Query): this;
 
+        /**
+         * Sets knn on the request body.
+         *
+         * @param {KNN|KNN[]} knn
+         */
+        kNN(knn: KNN | KNN[]): this;
+
         /**
          * Sets aggregation on the request body.
          * Alias for method `aggregation`
@@ -3141,7 +3148,7 @@ declare namespace esb {
 
         /**
          * Sets the script used to compute the score of documents returned by the query.
-         * 
+         *
          * @param {Script} script A valid `Script` object
          */
         script(script: Script): this;
@@ -3761,6 +3768,84 @@ declare namespace esb {
         spanQry?: SpanQueryBase
     ): SpanFieldMaskingQuery;
 
+    /**
+     * Knn performs k-nearest neighbor (KNN) searches.
+     * This class allows configuring the KNN search with various parameters such as field, query vector,
+     * number of nearest neighbors (k), number of candidates, boost factor, and similarity metric.
+     *
+     * NOTE: Only available in Elasticsearch v8.0+
+     */
+    export class KNN {
+        /**
+         * Creates an instance of Knn, initializing the internal state for the k-NN search.
+         *
+         * @param {string} field - (Optional) The field against which to perform the k-NN search.
+         * @param {number} k - (Optional) The number of nearest neighbors to retrieve.
+         * @param {number} numCandidates - (Optional) The number of candidate neighbors to consider during the search.
+         * @throws {Error} If the number of candidates (numCandidates) is less than the number of neighbors (k).
+         */
+        constructor(field: string, k: number, numCandidates: number);
+
+        /**
+         * Sets the query vector for the KNN search, an array of numbers representing the reference point.
+         *
+         * @param {number[]} vector
+         */
+        queryVector(vector: number[]): this;
+
+        /**
+         * Sets the query vector builder for the k-NN search.
+         * This method configures a query vector builder using a specified model ID and model text.
+         * Note that either a direct query vector or a query vector builder can be provided, but not both.
+         *
+         * @param {string} modelId - The ID of the model used for generating the query vector.
+         * @param {string} modelText - The text input based on which the query vector is generated.
+         * @returns {KNN} Returns the instance of Knn for method chaining.
+         * @throws {Error} If both query_vector_builder and query_vector are provided.
+         */
+        queryVectorBuilder(modelId: string, modelText: string): this;
+
+        /**
+         * Adds one or more filter queries to the k-NN search.
+         * This method is designed to apply filters to the k-NN search. It accepts either a single
+         * query or an array of queries. Each query acts as a filter, refining the search results
+         * according to the specified conditions. These queries must be instances of the `Query` class.
+         *
+         * @param {Query|Query[]} queries - A single `Query` instance or an array of `Query` instances for filtering.
+         * @returns {KNN} Returns `this` to allow method chaining.
+         * @throws {TypeError} If any of the provided queries is not an instance of `Query`.
+         */
+        filter(queries: Query | Query[]): this;
+
+        /**
+         * Applies a boost factor to the query to influence the relevance score of returned documents.
+         *
+         * @param {number} boost
+         */
+        boost(boost: number): this;
+
+        /**
+         * Sets the similarity metric used in the KNN algorithm to calculate similarity.
+         *
+         * @param {number} similarity
+         */
+        similarity(similarity: number): this;
+
+        /**
+         * Override default `toJSON` to return DSL representation for the `query`
+         *
+         * @override
+         */
+        toJSON(): object;
+    }
+
+    /**
+     * Factory function to instantiate a new Knn object.
+     *
+     * @returns {KNN}
+     */
+    export function kNN(field: string, k: number, numCandidates: number): KNN;
+
     /**
      * Base class implementation for all aggregation types.
      *
@@ -3913,9 +3998,9 @@ declare namespace esb {
     /**
      * A single-value metrics aggregation that computes the weighted average of numeric values that are extracted from the aggregated documents.
      * These values can be extracted either from specific numeric fields in the documents.
-     * 
+     *
      * [Elasticsearch reference](https://www.elastic.co/guide/en/elasticsearch/reference/current/search-aggregations-metrics-weight-avg-aggregation.html)
-     * 
+     *
      * Added in Elasticsearch v6.4.0
      * [Release notes](https://www.elastic.co/guide/en/elasticsearch/reference/6.4/release-notes-6.4.0.html)
      *
@@ -3929,7 +4014,7 @@ declare namespace esb {
 
         /**
          * Sets the value
-         * 
+         *
          * @param {string | Script} value Field name or script to be used as the value
          * @param {number=} missing Sets the missing parameter which defines how documents
          * that are missing a value should be treated.
@@ -3939,7 +4024,7 @@ declare namespace esb {
 
         /**
          * Sets the weight
-         * 
+         *
          * @param {string | Script} weight Field name or script to be used as the weight
          * @param {number=} missing Sets the missing parameter which defines how documents
          * that are missing a value should be treated.
@@ -3969,9 +4054,9 @@ declare namespace esb {
     /**
      * A single-value metrics aggregation that computes the weighted average of numeric values that are extracted from the aggregated documents.
      * These values can be extracted either from specific numeric fields in the documents.
-     * 
+     *
      * [Elasticsearch reference](https://www.elastic.co/guide/en/elasticsearch/reference/current/search-aggregations-metrics-weight-avg-aggregation.html)
-     * 
+     *
      * Added in Elasticsearch v6.4.0
      * [Release notes](https://www.elastic.co/guide/en/elasticsearch/reference/6.4/release-notes-6.4.0.html)
      *
@@ -8922,15 +9007,15 @@ declare namespace esb {
 
         /**
          * Sets the type of the runtime field.
-         * 
+         *
          * @param {string} type One of `boolean`, `composite`, `date`, `double`, `geo_point`, `ip`, `keyword`, `long`, `lookup`.
          * @returns {void}
          */
         type(type: 'boolean' | 'composite' | 'date' | 'double' | 'geo_point' | 'ip' | 'keyword' | 'long' | 'lookup'): void;
 
         /**
          * Sets the source of the script.
-         * 
+         *
          * @param {string} script
          * @returns {void}
          */
diff --git a/src/index.js b/src/index.js
@@ -15,6 +15,7 @@ const {
     RuntimeField,
     SearchTemplate,
     Query,
+    KNN,
     util: { constructorWrapper }
 } = require('./core');
 
@@ -343,6 +344,13 @@ exports.spanWithinQuery = constructorWrapper(SpanWithinQuery);
 
 exports.SpanFieldMaskingQuery = SpanFieldMaskingQuery;
 exports.spanFieldMaskingQuery = constructorWrapper(SpanFieldMaskingQuery);
+
+/* ============ ============ ============ */
+/* ======== KNN ======== */
+/* ============ ============ ============ */
+exports.KNN = KNN;
+exports.kNN = constructorWrapper(KNN);
+
 /* ============ ============ ============ */
 /* ======== Metrics Aggregations ======== */
 /* ============ ============ ============ */
diff --git a/test/core-test/knn.test.js b/test/core-test/knn.test.js
diff --git a/test/core-test/request-body-search.test.js b/test/core-test/request-body-search.test.js
