-
Notifications
You must be signed in to change notification settings - Fork 76
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
kNN search was added to Elasticsearch in v8.0 Co-authored-by: kennylindahl <[email protected]> Co-authored-by: Andreas Franzon <[email protected]>
- Loading branch information
3 people
authored
May 5, 2024
1 parent
8f73d34
commit 78de179
Showing
7 changed files
with
485 additions
and
11 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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; |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.